mysql数据链接更新
This commit is contained in:
+104
-114
@@ -3,18 +3,20 @@ import pandas as pd
|
||||
from datetime import datetime
|
||||
import time
|
||||
import pymysql
|
||||
from utils.mysql_agent import MySQLAgent
|
||||
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 = "test_db_" + datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
cls.test_table = "test_table_" + datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
# 创建唯一的测试数据库和表名（避免冲突）
|
||||
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,
|
||||
@@ -32,21 +34,19 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
'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'],
|
||||
@@ -54,7 +54,6 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
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}")
|
||||
@@ -66,21 +65,14 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
|
||||
@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(
|
||||
host=cls.base_config['host'],
|
||||
port=cls.base_config['port'],
|
||||
user=cls.base_config['user'],
|
||||
password=cls.base_config['password'],
|
||||
charset='utf8mb4'
|
||||
)
|
||||
|
||||
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}")
|
||||
@@ -88,22 +80,24 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
finally:
|
||||
temp_conn.close()
|
||||
|
||||
def test_01_connection(self):
|
||||
def test_connection(self):
|
||||
"""测试数据库连接"""
|
||||
version = self.db.query_to_df("SELECT VERSION() as version")
|
||||
self.assertIsNotNone(version)
|
||||
print(f"\nDatabase version: {version['version'].iloc[0]}")
|
||||
print(f"Running on: {platform.system()} {platform.release()}")
|
||||
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_02_query_to_df(self):
|
||||
def test_query_to_df(self):
|
||||
"""测试查询返回DataFrame"""
|
||||
df = self.db.query_to_df(f"SELECT * FROM {self.test_table} WHERE id > %s", (1,))
|
||||
self.assertEqual(len(df), 2)
|
||||
df = self.db.query_to_df(
|
||||
f"SELECT * FROM {self.test_table} WHERE id > %s",
|
||||
params=(1,)
|
||||
)
|
||||
self.assertIsInstance(df, pd.DataFrame)
|
||||
print("\nQuery result sample:")
|
||||
print(df.head())
|
||||
self.assertEqual(len(df), 2) # id>1 的数据有2条
|
||||
self.assertIn('name', df.columns)
|
||||
|
||||
def test_03_insert_from_df(self):
|
||||
def test_insert_from_df(self):
|
||||
"""测试DataFrame插入"""
|
||||
new_data = pd.DataFrame({
|
||||
'id': [4, 5],
|
||||
@@ -112,55 +106,55 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
'created_at': pd.to_datetime(['2023-01-04', '2023-01-05'])
|
||||
})
|
||||
|
||||
rows = self.db.insert_from_df(self.test_table, new_data)
|
||||
self.assertEqual(rows, 2)
|
||||
inserted_rows = self.db.insert_from_df(self.test_table, new_data)
|
||||
self.assertEqual(inserted_rows, 2)
|
||||
|
||||
# 验证数据
|
||||
df = self.db.query_to_df(f"SELECT * FROM {self.test_table} WHERE id >= 4")
|
||||
self.assertEqual(len(df), 2)
|
||||
self.assertEqual(df['name'].tolist(), ['Test4', 'Test5'])
|
||||
# 验证插入结果
|
||||
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_04_update_from_df(self):
|
||||
def test_update_from_df(self):
|
||||
"""测试DataFrame更新"""
|
||||
update_data = pd.DataFrame({
|
||||
'id': [1, 2],
|
||||
'name': ['Updated1', 'Updated2']
|
||||
})
|
||||
|
||||
rows = self.db.update_from_df(self.test_table, update_data, 'id')
|
||||
self.assertGreaterEqual(rows, 2)
|
||||
updated_rows = self.db.update_from_df(self.test_table, update_data, 'id')
|
||||
self.assertGreaterEqual(updated_rows, 2)
|
||||
|
||||
# 验证更新
|
||||
df = self.db.query_to_df(f"SELECT name FROM {self.test_table} WHERE id IN (1,2)")
|
||||
self.assertIn('Updated1', df['name'].values)
|
||||
self.assertIn('Updated2', df['name'].values)
|
||||
# 验证更新结果
|
||||
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_05_transaction(self):
|
||||
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")
|
||||
|
||||
# 验证事务内修改
|
||||
cursor.execute(f"SELECT value FROM {self.test_table} WHERE id = 1")
|
||||
self.assertEqual(cursor.fetchone()['value'], 99.9)
|
||||
|
||||
self.db.commit_transaction(conn)
|
||||
except Exception:
|
||||
self.db.rollback_transaction(conn)
|
||||
raise
|
||||
|
||||
# 验证提交后的修改
|
||||
df = self.db.query_to_df(f"SELECT value FROM {self.test_table} WHERE id IN (1,2)")
|
||||
self.assertIn(99.9, df['value'].values)
|
||||
self.assertIn(88.8, df['value'].values)
|
||||
# 验证事务提交结果
|
||||
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_06_large_data(self):
|
||||
"""测试大数据量操作"""
|
||||
# 生成测试数据
|
||||
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)],
|
||||
@@ -168,59 +162,55 @@ class TestMySQLAgent(unittest.TestCase):
|
||||
'created_at': pd.date_range('2023-01-01', periods=1000)
|
||||
})
|
||||
|
||||
# Windows平台使用更小的批次
|
||||
# 根据平台自动调整批次大小
|
||||
chunk_size = 100 if platform.system() == 'Windows' else 500
|
||||
|
||||
start_time = time.time()
|
||||
rows = self.db.insert_from_df(self.test_table, large_data, chunk_size=chunk_size)
|
||||
inserted_rows = self.db.insert_from_df(
|
||||
self.test_table,
|
||||
large_data,
|
||||
chunk_size=chunk_size
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
self.assertEqual(rows, 1000)
|
||||
print(f"\nInserted 1000 rows in {elapsed:.2f}s (chunk_size={chunk_size})")
|
||||
self.assertEqual(inserted_rows, 1000)
|
||||
print(f"插入1000行数据耗时: {elapsed:.2f}秒 (批次大小: {chunk_size})")
|
||||
|
||||
# 验证数据
|
||||
df = self.db.query_to_df(f"SELECT COUNT(*) as cnt FROM {self.test_table} WHERE id >= 1000")
|
||||
self.assertEqual(df['cnt'].iloc[0], 1000)
|
||||
|
||||
def test_07_concurrent_access(self):
|
||||
def test_concurrent_access(self):
|
||||
"""测试并发访问"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
def worker(i):
|
||||
df = self.db.query_to_df(f"SELECT * FROM {self.test_table} WHERE id = %s", (i % 5 + 1,))
|
||||
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(worker, range(100)))
|
||||
|
||||
results = list(executor.map(query_worker, range(100)))
|
||||
elapsed = time.time() - start_time
|
||||
self.assertEqual(sum(results), 100)
|
||||
print(f"\nCompleted 100 concurrent queries in {elapsed:.2f}s")
|
||||
|
||||
self.assertEqual(sum(results), 100) # 每次查询应返回1行
|
||||
print(f"100次并发查询耗时: {elapsed:.2f}秒")
|
||||
|
||||
|
||||
class TestPlatformSpecific(unittest.TestCase):
|
||||
"""平台特定功能测试"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""创建临时测试数据库"""
|
||||
cls.test_db_name = "test_db_platform_" + datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
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',
|
||||
'max_connections': 10
|
||||
'password': '123123'
|
||||
}
|
||||
|
||||
# 创建数据库
|
||||
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'
|
||||
)
|
||||
|
||||
# 创建测试数据库
|
||||
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}")
|
||||
@@ -230,15 +220,8 @@ class TestPlatformSpecific(unittest.TestCase):
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(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'
|
||||
)
|
||||
|
||||
"""清理测试数据库"""
|
||||
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}")
|
||||
@@ -249,42 +232,49 @@ class TestPlatformSpecific(unittest.TestCase):
|
||||
def test_windows_timeout(self):
|
||||
"""测试Windows平台超时处理"""
|
||||
if platform.system() != 'Windows':
|
||||
self.skipTest("Only runs on Windows")
|
||||
self.skipTest("仅在Windows平台运行")
|
||||
|
||||
config = {
|
||||
**self.base_config,
|
||||
'database': self.test_db_name,
|
||||
'connect_timeout': 1,
|
||||
'read_timeout': 1
|
||||
'read_timeout': 1,
|
||||
'write_timeout': 1
|
||||
}
|
||||
|
||||
db = MySQLAgent(config)
|
||||
|
||||
# 测试短超时查询
|
||||
start_time = time.time()
|
||||
try:
|
||||
db.query_to_df("SELECT SLEEP(2)")
|
||||
self.fail("Should have timed out")
|
||||
except Exception as e:
|
||||
self.assertIn("timed out", str(e))
|
||||
print(f"\nWindows timeout test: {str(e)}")
|
||||
# 执行会超时的查询(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
|
||||
|
||||
def test_macos_ssl(self):
|
||||
"""测试macOS SSL连接"""
|
||||
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("Only runs on macOS")
|
||||
self.skipTest("仅在macOS平台运行")
|
||||
|
||||
config = {
|
||||
**self.base_config,
|
||||
'database': self.test_db_name,
|
||||
'ssl': {'ca': '/usr/local/etc/openssl/cert.pem'}
|
||||
}
|
||||
|
||||
db = MySQLAgent(config)
|
||||
version = db.query_to_df("SELECT VERSION() as version")
|
||||
self.assertIsNotNone(version)
|
||||
print(f"\nmacOS SSL connection successful: {version['version'].iloc[0]}")
|
||||
version_df = db.query_to_df("SELECT VERSION() as version")
|
||||
self.assertIsNotNone(version_df)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user