RagFlow_Plugs/test/test_mysql_conn.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MySQL连接工具类测试脚本
"""
import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.mysql_conn import get_mysql_conn

def test_mysql_conn():
    """
    测试MySQL连接工具类
    """
    print("测试MySQL连接工具类...")
    
    try:
        # 获取MySQL连接管理器实例
        mysql_conn = get_mysql_conn(
            host="localhost",
            port=3306,
            user="root",
            password="password",
            database="test_db",
            pool_size=3
        )
        
        print("✓ 成功获取MySQL连接管理器实例")
        
        # 测试创建表
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS test_users (
            id INT AUTO_INCREMENT PRIMARY KEY,
            name VARCHAR(50) NOT NULL,
            email VARCHAR(100) NOT NULL UNIQUE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """
        
        result = mysql_conn.execute(create_table_sql)
        print("✓ 成功创建测试表")
        
        # 测试插入数据
        insert_sql = "INSERT INTO test_users (name, email) VALUES (%s, %s)"
        insert_params = ("测试用户", "test@example.com")
        
        row_count = mysql_conn.execute(insert_sql, insert_params)
        print(f"✓ 成功插入 {row_count} 条数据")
        
        # 测试查询数据
        select_sql = "SELECT * FROM test_users WHERE name = %s"
        select_params = ("测试用户",)
        
        user = mysql_conn.fetch_one(select_sql, select_params)
        if user:
            print(f"✓ 成功查询数据: {user}")
        else:
            print("✗ 查询数据失败")
        
        # 测试批量插入
        bulk_insert_sql = "INSERT INTO test_users (name, email) VALUES (%s, %s)"
        bulk_params = [
            ("批量用户1", "batch1@example.com"),
            ("批量用户2", "batch2@example.com"),
            ("批量用户3", "batch3@example.com")
        ]
        
        bulk_row_count = mysql_conn.bulk_insert(bulk_insert_sql, bulk_params)
        print(f"✓ 成功批量插入 {bulk_row_count} 条数据")
        
        # 测试查询所有数据
        select_all_sql = "SELECT * FROM test_users"
        all_users = mysql_conn.fetch_all(select_all_sql)
        print(f"✓ 成功查询所有数据，共 {len(all_users)} 条")
        
        # 测试更新数据
        update_sql = "UPDATE test_users SET name = %s WHERE id = %s"
        update_params = ("更新后的测试用户", user["id"])
        
        update_row_count = mysql_conn.execute(update_sql, update_params)
        print(f"✓ 成功更新 {update_row_count} 条数据")
        
        # 测试删除数据
        delete_sql = "DELETE FROM test_users WHERE id = %s"
        delete_params = (user["id"],)
        
        delete_row_count = mysql_conn.execute(delete_sql, delete_params)
        print(f"✓ 成功删除 {delete_row_count} 条数据")
        
        # 测试事务
        print("测试事务处理...")
        conn, cursor = mysql_conn.begin_transaction()
        try:
            # 在事务中执行多个操作
            cursor.execute("INSERT INTO test_users (name, email) VALUES (%s, %s)", ("事务用户1", "transaction1@example.com"))
            cursor.execute("INSERT INTO test_users (name, email) VALUES (%s, %s)", ("事务用户2", "transaction2@example.com"))
            mysql_conn.commit_transaction(conn, cursor)
            print("✓ 事务提交成功")
        except Exception as e:
            mysql_conn.rollback_transaction(conn, cursor)
            print(f"✗ 事务回滚: {e}")
        
        # 清理测试数据
        drop_table_sql = "DROP TABLE IF EXISTS test_users"
        mysql_conn.execute(drop_table_sql)
        print("✓ 成功清理测试表")
        
        # 关闭连接池
        mysql_conn.close()
        print("✓ 成功关闭连接池")
        
        print("\n🎉 所有测试通过！MySQL连接工具类工作正常。")
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True

if __name__ == "__main__":
    test_mysql_conn()