import unittest import pandas as pd from datetime import datetime import time import pymysql import platform from concurrent.futures import ThreadPoolExecutor from utils.mysql_agent import MySQLAgent class TestMySQLAgent(unittest.TestCase): @classmethod def setUpClass(cls): """初始化测试环境和测试表""" # 创建唯一的测试数据库和表名（避免冲突） cls.test_db_name = f"test_db_{datetime.now().strftime('%Y%m%d%H%M%S')}" cls.test_table = f"test_table_{datetime.now().strftime('%Y%m%d%H%M%S')}" # 基础配置（根据实际环境修改） cls.base_config = { 'host': 'localhost', 'port': 3306, 'user': 'root', 'password': '123123', 'max_connections': 10 } # 创建测试数据库 cls._create_test_database() # 初始化数据库连接 cls.db = MySQLAgent({ **cls.base_config, 'database': cls.test_db_name }) # 创建测试表并插入初始数据 test_data = pd.DataFrame({ 'id': [1, 2, 3], 'name': ['Test1', 'Test2', 'Test3'], 'value': [10.5, 20.3, 30.8], 'created_at': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']) }) cls.db.create_table_from_df(cls.test_table, test_data, primary_key='id') cls.db.insert_from_df(cls.test_table, test_data) @classmethod def _create_test_database(cls): """创建测试数据库""" temp_conn = pymysql.connect( host=cls.base_config['host'], port=cls.base_config['port'], user=cls.base_config['user'], password=cls.base_config['password'], charset='utf8mb4' ) try: with temp_conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE IF NOT EXISTS {cls.test_db_name}") cursor.execute(f"USE {cls.test_db_name}") cursor.execute("SET GLOBAL max_connections = 100") temp_conn.commit() finally: temp_conn.close() @classmethod def tearDownClass(cls): """清理测试环境""" if hasattr(cls, 'db') and cls.db: # 删除测试表 if cls.db.table_exists(cls.test_table): cls.db.drop_table(cls.test_table) # 删除测试数据库 temp_conn = pymysql.connect(**cls.base_config, charset='utf8mb4') try: with temp_conn.cursor() as cursor: cursor.execute(f"DROP DATABASE IF EXISTS {cls.test_db_name}") temp_conn.commit() finally: temp_conn.close() def test_connection(self): """测试数据库连接""" version_df = self.db.query_to_df("SELECT VERSION() as version") self.assertIsNotNone(version_df) self.assertEqual(len(version_df), 1) print(f"数据库版本: {version_df['version'].iloc[0]}") def test_query_to_df(self): """测试查询返回DataFrame""" df = self.db.query_to_df( f"SELECT * FROM {self.test_table} WHERE id > %s", params=(1,) ) self.assertIsInstance(df, pd.DataFrame) self.assertEqual(len(df), 2) # id>1 的数据有2条 self.assertIn('name', df.columns) def test_insert_from_df(self): """测试DataFrame插入""" new_data = pd.DataFrame({ 'id': [4, 5], 'name': ['Test4', 'Test5'], 'value': [40.1, 50.2], 'created_at': pd.to_datetime(['2023-01-04', '2023-01-05']) }) inserted_rows = self.db.insert_from_df(self.test_table, new_data) self.assertEqual(inserted_rows, 2) # 验证插入结果 result_df = self.db.query_to_df( f"SELECT name FROM {self.test_table} WHERE id IN (4,5)" ) self.assertEqual(result_df['name'].tolist(), ['Test4', 'Test5']) def test_update_from_df(self): """测试DataFrame更新""" update_data = pd.DataFrame({ 'id': [1, 2], 'name': ['Updated1', 'Updated2'] }) updated_rows = self.db.update_from_df(self.test_table, update_data, 'id') self.assertGreaterEqual(updated_rows, 2) # 验证更新结果 result_df = self.db.query_to_df( f"SELECT name FROM {self.test_table} WHERE id IN (1,2)" ) self.assertIn('Updated1', result_df['name'].values) self.assertIn('Updated2', result_df['name'].values) def test_transaction(self): """测试事务处理""" conn = self.db.begin_transaction() try: # 执行事务内操作 cursor = conn.cursor() cursor.execute(f"UPDATE {self.test_table} SET value = 99.9 WHERE id = 1") cursor.execute(f"UPDATE {self.test_table} SET value = 88.8 WHERE id = 2") self.db.commit_transaction(conn) except Exception: self.db.rollback_transaction(conn) raise # 验证事务提交结果 result_df = self.db.query_to_df( f"SELECT value FROM {self.test_table} WHERE id IN (1,2)" ) self.assertIn(99.9, result_df['value'].values) self.assertIn(88.8, result_df['value'].values) def test_large_data_insert(self): """测试大数据量插入""" # 生成1000行测试数据 large_data = pd.DataFrame({ 'id': range(1000, 2000), 'name': [f"Item_{i}" for i in range(1000, 2000)], 'value': [i * 0.1 for i in range(1000, 2000)], 'created_at': pd.date_range('2023-01-01', periods=1000) }) # 根据平台自动调整批次大小 chunk_size = 100 if platform.system() == 'Windows' else 500 start_time = time.time() inserted_rows = self.db.insert_from_df( self.test_table, large_data, chunk_size=chunk_size ) elapsed = time.time() - start_time self.assertEqual(inserted_rows, 1000) print(f"插入1000行数据耗时: {elapsed:.2f}秒 (批次大小: {chunk_size})") def test_concurrent_access(self): """测试并发访问""" def query_worker(i): """并发查询工作函数""" df = self.db.query_to_df( f"SELECT * FROM {self.test_table} WHERE id = %s", params=(i % 3 + 1,) # 查询id=1,2,3循环 ) return len(df) # 20个线程执行100次查询 start_time = time.time() with ThreadPoolExecutor(max_workers=20) as executor: results = list(executor.map(query_worker, range(100))) elapsed = time.time() - start_time self.assertEqual(sum(results), 100) # 每次查询应返回1行 print(f"100次并发查询耗时: {elapsed:.2f}秒") class TestPlatformSpecific(unittest.TestCase): """平台特定功能测试""" @classmethod def setUpClass(cls): cls.test_db_name = f"test_platform_db_{datetime.now().strftime('%Y%m%d%H%M%S')}" cls.base_config = { 'host': 'localhost', 'port': 3306, 'user': 'root', 'password': '123123' } # 创建测试数据库 temp_conn = pymysql.connect(**cls.base_config, charset='utf8mb4') try: with temp_conn.cursor() as cursor: cursor.execute(f"CREATE DATABASE IF NOT EXISTS {cls.test_db_name}") temp_conn.commit() finally: temp_conn.close() @classmethod def tearDownClass(cls): """清理测试数据库""" temp_conn = pymysql.connect(**cls.base_config, charset='utf8mb4') try: with temp_conn.cursor() as cursor: cursor.execute(f"DROP DATABASE IF EXISTS {cls.test_db_name}") temp_conn.commit() finally: temp_conn.close() def test_windows_timeout(self): """测试Windows平台超时处理""" if platform.system() != 'Windows': self.skipTest("仅在Windows平台运行") config = { **self.base_config, 'database': self.test_db_name, 'connect_timeout': 1, 'read_timeout': 1, 'write_timeout': 1 } db = MySQLAgent(config) # 执行会超时的查询(SLEEP(2)超过1秒超时设置) with self.assertRaises((pymysql.OperationalError, TimeoutError)) as ctx: try: db.query_to_df("SELECT SLEEP(2)") except Exception as e: # 提取底层异常信息（可能被包装） while hasattr(e, 'args') and len(e.args) > 0 and isinstance(e.args[0], Exception): e = e.args[0] raise e error_msg = str(ctx.exception) self.assertTrue( "timed out" in error_msg or "timeout" in error_msg or "HY000" in error_msg, # MySQL超时错误码 f"未检测到超时异常，实际异常: {error_msg}" ) def test_macos_ssl_connection(self): """测试macOS平台SSL连接""" if platform.system() != 'Darwin': self.skipTest("仅在macOS平台运行") config = { **self.base_config, 'database': self.test_db_name, 'ssl': {'ca': '/usr/local/etc/openssl/cert.pem'} } db = MySQLAgent(config) version_df = db.query_to_df("SELECT VERSION() as version") self.assertIsNotNone(version_df) if __name__ == '__main__': unittest.main(verbosity=2)