#!/usr/bin/env python3
"""
测试优雅关闭功能

测试 CeleryChannel 的优雅关闭机制，确保在进程终止时能正确释放未完成的任务。
"""

import unittest
import time
import threading
from unittest.mock import Mock, MagicMock, patch
import sys
import os
from pathlib import Path

# Add the parent directory to the path to import sandai
sys.path.insert(0, str(Path(__file__).parent.parent))

from sandai.operator.channel import CeleryChannel


class MockMessage:
    """模拟 Celery 消息对象"""
    
    def __init__(self, task_id: str):
        self.task_id = task_id
        self.rejected = False
        self.acked = False
        self.properties = {'correlation_id': task_id}
        self.delivery_info = {'routing_key': f'normal.test-operator.v1'}
    
    def reject(self):
        """模拟 reject 操作"""
        self.rejected = True
        print(f"Message {self.task_id} rejected")
    
    def ack(self):
        """模拟 ack 操作"""
        self.acked = True
        print(f"Message {self.task_id} acknowledged")


class TestGracefulShutdown(unittest.TestCase):
    """测试优雅关闭功能"""

    def setUp(self):
        """设置测试环境"""
        # 模拟 Celery 相关的导入
        self.mock_celery_app = Mock()
        self.mock_connection = Mock()
        self.mock_consumer = Mock()
        
        # 创建 CeleryChannel 实例（不启动线程）
        with patch('sandai.operator.channel.Celery'), \
             patch('sandai.operator.env.get_celery_config') as mock_get_config, \
             patch('sandai.operator.env.get_default_celery_config') as mock_default_config:
            
            # 模拟配置
            mock_config = Mock()
            mock_config.broker_url = "redis://localhost:6379/0"
            mock_config.result_backend = "redis://localhost:6379/0"
            mock_config.task_serializer = "json"
            mock_config.result_serializer = "json"
            mock_config.accept_content = ["json"]
            
            mock_get_config.return_value = mock_config
            mock_default_config.return_value = mock_config
            
            self.channel = CeleryChannel(
                operator_name="test-operator",
                operator_version="v1",
                visibility_timeout=300,  # 5 分钟
                debug=True
            )
            
            # 阻止自动启动线程
            self.channel._stop_event.set()

    def test_visibility_timeout_default(self):
        """测试默认的 visibility_timeout 是否为 5 分钟"""
        with patch('sandai.operator.channel.Celery'), \
             patch('sandai.operator.env.get_celery_config') as mock_get_config, \
             patch('sandai.operator.env.get_default_celery_config') as mock_default_config:
            
            mock_config = Mock()
            mock_config.broker_url = "redis://localhost:6379/0"
            mock_config.result_backend = "redis://localhost:6379/0"
            mock_config.task_serializer = "json"
            mock_config.result_serializer = "json"
            mock_config.accept_content = ["json"]
            
            mock_get_config.return_value = mock_config
            mock_default_config.return_value = mock_config
            
            # 创建不指定 visibility_timeout 的 channel
            channel = CeleryChannel(
                operator_name="test-operator",
                operator_version="v1",
                debug=True
            )
            
            # 验证默认值是 300 秒（5 分钟）
            self.assertEqual(channel.visibility_timeout, 300)

    def test_graceful_shutdown_empty_process_map(self):
        """测试空进程映射的优雅关闭"""
        # 确保进程映射为空
        self.channel._current_process_map = {}
        
        # 执行优雅关闭
        self.channel.graceful_shutdown()
        
        # 验证停止事件被设置
        self.assertTrue(self.channel._stop_event.is_set())
        
        # 验证进程映射仍为空
        self.assertEqual(len(self.channel._current_process_map), 0)

    def test_graceful_shutdown_with_pending_tasks(self):
        """测试有待处理任务时的优雅关闭"""
        # 创建模拟的待处理任务
        mock_messages = []
        for i in range(3):
            task_id = f"test_task_{i}"
            message = MockMessage(task_id)
            mock_messages.append(message)
            
            # 添加到进程映射
            self.channel._current_process_map[task_id] = {
                'task_data': {'task_id': task_id, 'options': {}},
                'message': message,
                'celery_task_id': task_id
            }
        
        # 验证初始状态
        self.assertEqual(len(self.channel._current_process_map), 3)
        
        # 执行优雅关闭
        self.channel.graceful_shutdown()
        
        # 验证所有任务都被 reject
        for message in mock_messages:
            self.assertTrue(message.rejected, f"Message {message.task_id} should be rejected")
        
        # 验证进程映射被清空
        self.assertEqual(len(self.channel._current_process_map), 0)
        
        # 验证停止事件被设置
        self.assertTrue(self.channel._stop_event.is_set())

    def test_graceful_shutdown_with_error_handling(self):
        """测试优雅关闭时的错误处理"""
        # 创建一个会抛出异常的模拟消息
        class ErrorMessage(MockMessage):
            def reject(self):
                raise Exception("Reject failed")
        
        # 添加正常消息和错误消息
        normal_message = MockMessage("normal_task")
        error_message = ErrorMessage("error_task")
        
        self.channel._current_process_map = {
            "normal_task": {
                'task_data': {'task_id': 'normal_task'},
                'message': normal_message,
                'celery_task_id': 'normal_task'
            },
            "error_task": {
                'task_data': {'task_id': 'error_task'},
                'message': error_message,
                'celery_task_id': 'error_task'
            }
        }
        
        # 执行优雅关闭（不应该抛出异常）
        try:
            self.channel.graceful_shutdown()
        except Exception as e:
            self.fail(f"Graceful shutdown should handle errors gracefully, but got: {e}")
        
        # 验证正常消息被 reject
        self.assertTrue(normal_message.rejected)
        
        # 验证进程映射被清空（即使有错误）
        self.assertEqual(len(self.channel._current_process_map), 0)

    def test_close_calls_graceful_shutdown(self):
        """测试 close 方法会调用优雅关闭"""
        # 添加一个待处理任务
        message = MockMessage("test_task")
        self.channel._current_process_map["test_task"] = {
            'task_data': {'task_id': 'test_task'},
            'message': message,
            'celery_task_id': 'test_task'
        }
        
        # 重置停止事件（模拟正常运行状态）
        self.channel._stop_event.clear()
        
        # 调用 close
        self.channel.close()
        
        # 验证任务被 reject
        self.assertTrue(message.rejected)
        
        # 验证停止事件被设置
        self.assertTrue(self.channel._stop_event.is_set())

    def test_multiple_graceful_shutdowns(self):
        """测试多次调用优雅关闭的安全性"""
        # 添加一个待处理任务
        message = MockMessage("test_task")
        self.channel._current_process_map["test_task"] = {
            'task_data': {'task_id': 'test_task'},
            'message': message,
            'celery_task_id': 'test_task'
        }
        
        # 第一次优雅关闭
        self.channel.graceful_shutdown()
        self.assertTrue(message.rejected)
        
        # 重置消息状态
        message.rejected = False
        
        # 第二次优雅关闭（应该安全，不会重复 reject）
        self.channel.graceful_shutdown()
        
        # 验证消息没有被重复 reject（因为进程映射已清空）
        self.assertFalse(message.rejected)

    def test_concurrent_graceful_shutdown(self):
        """测试并发优雅关闭的线程安全性"""
        # 添加多个待处理任务
        messages = []
        for i in range(10):
            task_id = f"concurrent_task_{i}"
            message = MockMessage(task_id)
            messages.append(message)
            
            self.channel._current_process_map[task_id] = {
                'task_data': {'task_id': task_id},
                'message': message,
                'celery_task_id': task_id
            }
        
        # 创建多个线程同时执行优雅关闭
        threads = []
        for i in range(3):
            thread = threading.Thread(target=self.channel.graceful_shutdown)
            threads.append(thread)
        
        # 启动所有线程
        for thread in threads:
            thread.start()
        
        # 等待所有线程完成
        for thread in threads:
            thread.join(timeout=1.0)
        
        # 验证所有任务都被处理（每个任务只被 reject 一次）
        rejected_count = sum(1 for msg in messages if msg.rejected)
        self.assertEqual(rejected_count, 10, "All tasks should be rejected exactly once")
        
        # 验证进程映射被清空
        self.assertEqual(len(self.channel._current_process_map), 0)


def run_graceful_shutdown_tests():
    """运行优雅关闭测试"""
    # 创建测试套件
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()
    
    # 添加测试用例
    suite.addTests(loader.loadTestsFromTestCase(TestGracefulShutdown))
    
    # 运行测试
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    
    return result.wasSuccessful()


if __name__ == '__main__':
    # 运行测试
    success = run_graceful_shutdown_tests()
    
    if success:
        print("\n🎉 All graceful shutdown tests passed!")
        sys.exit(0)
    else:
        print("\n💥 Some graceful shutdown tests failed!")
        sys.exit(1)