LQAgentPlatform/test/system_trace_id_test.py

"""
系统Trace ID测试
验证trace_id在异步并发和队列中的正确传播
"""
import os
import sys
# Add the parent directory (LQAgentPlatform) to sys.path so we can import foundation
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)
import asyncio
import time
from foundation.trace.trace_context import TraceContext, auto_trace
from foundation.logger.loggering import server_logger as logger


class TraceIDTest:
    """Trace ID测试类"""

    @staticmethod
    async def test_basic_context():
        """测试基础上下文功能"""
        print("\n=== 测试1: 基础上下文功能 ===")

        # 设置trace_id
        trace_id = TraceContext.generate_trace_id()
        TraceContext.set_trace_id(trace_id)

        logger.info("测试基础日志，应该包含trace_id")
        logger.info(f"手动设置的trace_id: {trace_id}")
        logger.info(f"自动获取的trace_id: {TraceContext.get_trace_id()}")

        assert TraceContext.get_trace_id() == trace_id, "trace_id设置失败"
        print("[PASS] 基础上下文功能测试通过")

    @staticmethod
    async def test_async_propagation():
        """测试异步并发传播"""
        print("\n=== 测试2: 异步并发传播 ===")

        # 设置主trace_id
        main_trace = "main-async-test"
        TraceContext.set_trace_id(main_trace)

        logger.info("主异步任务开始")

        async def concurrent_task(task_id: int):
            """并发任务"""
            current_trace = TraceContext.get_trace_id()
            logger.info(f"并发任务 {task_id} 获取到的trace_id: {current_trace}")

            # 在异步任务中修改trace_id，应该不影响其他任务
            new_trace = f"{main_trace}-subtask-{task_id}"
            TraceContext.set_trace_id(new_trace)

            await asyncio.sleep(0.1)
            logger.info(f"并发任务 {task_id} 修改后的trace_id: {new_trace}")

            return current_trace

        # 创建并发任务
        tasks = [concurrent_task(i) for i in range(3)]
        results = await asyncio.gather(*tasks)

        # 验证所有任务都继承到了主trace_id
        for i, result in enumerate(results):
            assert result == main_trace, f"任务 {i} 没有继承主trace_id"

        # 验证主任务trace_id不受影响
        assert TraceContext.get_trace_id() == main_trace, "主trace_id被并发任务污染"

        logger.info("主异步任务完成")
        print("[PASS] 异步并发传播测试通过")

    @staticmethod
    @auto_trace('callback_task_id')
    async def test_decorator_auto_trace(callback_task_id: str):
        """测试装饰器自动trace"""
        print(f"\n=== 测试3: 装饰器自动trace ===")

        # 不需要手动设置trace_id，装饰器会自动处理
        current_trace = TraceContext.get_trace_id()
        logger.info("装饰器自动设置的日志")

        assert current_trace == callback_task_id, "装饰器没有正确设置trace_id"

        # 测试装饰器在异步并发中的表现
        async def nested_task():
            nested_trace = TraceContext.get_trace_id()
            logger.info("嵌套异步任务")
            return nested_trace

        nested_result = await nested_task()
        assert nested_result == callback_task_id, "嵌套任务没有继承装饰器设置的trace_id"

        print(f"[PASS] 装饰器自动trace测试通过，trace_id: {callback_task_id}")

    @staticmethod
    async def test_context_manager():
        """测试上下文管理器"""
        print("\n=== 测试4: 上下文管理器 ===")

        original_trace = TraceContext.get_trace_id()
        logger.info(f"原始trace_id: {original_trace}")

        # 使用上下文管理器临时设置trace_id
        temp_trace = "temporary-trace"
        with TraceContext.with_trace_context(temp_trace) as ctx:
            logger.info("上下文管理器内的日志")
            current_trace = TraceContext.get_trace_id()
            assert current_trace == temp_trace, "上下文管理器没有正确设置trace_id"

        # 退出上下文后应该恢复原始trace_id
        restored_trace = TraceContext.get_trace_id()
        logger.info(f"恢复后的trace_id: {restored_trace}")
        assert restored_trace == original_trace, "上下文管理器没有正确恢复trace_id"

        print("[PASS] 上下文管理器测试通过")

    @staticmethod
    def test_celery_task_simulation():
        """测试Celery任务trace_id模拟"""
        print("\n=== 测试5: Celery任务trace_id模拟 ===")

        # 模拟提交Celery任务前的trace_id设置
        submit_trace = "celery-submit-test"
        TraceContext.set_trace_id(submit_trace)

        logger.info("准备提交Celery任务")

        # 模拟Celery任务执行
        def simulate_celery_task_execution(file_info: dict, _system_trace_id=None):
            """模拟Celery任务执行"""
            if _system_trace_id:
                TraceContext.set_trace_id(_system_trace_id)

            current_trace = TraceContext.get_trace_id()
            logger.info("Celery任务执行中")
            logger.info(f"文件ID: {file_info.get('file_id')}")

            return current_trace

        # 提交任务（模拟）
        file_info = {'file_id': 'test-file-123'}
        extracted_trace = TraceContext.get_trace_id()

        # 执行任务
        task_trace = simulate_celery_task_execution(
            file_info,
            _system_trace_id=extracted_trace
        )

        assert task_trace == submit_trace, "Celery任务没有正确获取到trace_id"

        print("[PASS] Celery任务trace_id模拟测试通过")


async def run_all_tests():
    """运行所有测试"""
    print("开始运行系统Trace ID测试...\n")

    try:
        # 测试1: 基础上下文功能
        await TraceIDTest.test_basic_context()

        # 测试2: 异步并发传播
        await TraceIDTest.test_async_propagation()

        # 测试3: 装饰器自动trace
        await TraceIDTest.test_decorator_auto_trace("decorator-test-123")

        # 测试4: 上下文管理器
        await TraceIDTest.test_context_manager()

        # 测试5: Celery任务模拟
        TraceIDTest.test_celery_task_simulation()

        print("\n[SUCCESS] 所有测试通过！系统Trace ID机制工作正常")
        return True

    except Exception as e:
        print(f"\n[FAIL] 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    # 运行测试
    success = asyncio.run(run_all_tests())
    exit(0 if success else 1)