TermControlAgent/tests/test_websocket_complete.py

��#!/usr/bin/env python3
"""
WebSocket�[te�R��KmՋ
g�R
����@b	gWebSocket�~�N�vck8^�]\O
"""
import asyncio
import json
import time
from datetime import datetime
from app.core.websocket.manager import websocket_manager
from app.core.websocket.channel import ChannelMessage
from app.utils.structured_log import get_structured_logger, LogLevel

logger = get_structured_logger(__name__, LogLevel.DEBUG)


class WebSocketCompleteTest:
    """WebSocket�[te�R��KmՋ{|"""
    
    def __init__(self):
        self.test_results = []
        self.test_client_name = "test_complete_client"
        self.test_url = "ws://localhost:8080/ws"
        
    async def run_all_tests(self):
        """ЏL�@b	gKmՋ"""
        logger.info("_�YWebSocket�[te�R��KmՋ")
        
        try:
            await self.test_create_client_and_channels()
            await self.test_connect_client()
            await self.test_send_messages()
            await self.test_receive_message_processing()
            await self.test_heartbeat_functionality()
            await self.test_priority_control()
            await self.test_disconnect_client()
            await self.test_cleanup()
            self.print_test_results()
            
        except Exception as e:
            logger.error(f"KmՋǏz-N�Su_8^: {e}")
            self.add_test_result("teSOKmՋ", False, f"_8^: {e}")
            self.print_test_results()
    
    async def test_create_client_and_channels(self):
        """KmՋ1�R�^�[7b�z�TChannel"""
        logger.info("KmՋ1�R�^�[7b�z�TChannel")
        
        try:
            client = await websocket_manager.create_client(
                name=self.test_client_name,
                url=self.test_url,
                heartbeat_interval=5
            )
            
            if client:
                self.add_test_result("R�^�[7b�z", True, "�[7b�zR�^b�R")
            else:
                self.add_test_result("R�^�[7b�z", False, "�[7b�zR�^1Y%�")
                return
            
            channels = websocket_manager.get_client_channels(self.test_client_name)
            expected_channels = [
                f"{self.test_client_name}_send",
                f"{self.test_client_name}_heartbeat", 
                f"{self.test_client_name}_receive"
            ]
            
            channel_names = list(channels.keys())
            if all(name in channel_names for name in expected_channels):
                self.add_test_result("R�^Channel", True, f"@b	gChannelR�^b�R: {channel_names}")
            else:
                self.add_test_result("R�^Channel", False, f"ChannelR�^
N�[te: {channel_names}")
                
        except Exception as e:
            self.add_test_result("R�^�[7b�z�TChannel", False, f"_8^: {e}")
    
    async def test_connect_client(self):
        """KmՋ2�ޏ�c�[7b�z"""
        logger.info("KmՋ2�ޏ�c�[7b�z")
        
        try:
            success = await websocket_manager.connect_client(self.test_client_name)
            
            if success:
                self.add_test_result("ޏ�c�[7b�z", True, "�[7b�zޏ�cb�R")
                
                client = websocket_manager.get_client(self.test_client_name)
                if client and client.is_connected:
                    self.add_test_result("�[7b�z�r
`", True, "�[7b�z�]ޏ�c")
                else:
                    self.add_test_result("�[7b�z�r
`", False, "�[7b�z*gޏ�c")
            else:
                self.add_test_result("ޏ�c�[7b�z", False, "�[7b�zޏ�c1Y%�")
                
        except Exception as e:
            self.add_test_result("ޏ�c�[7b�z", False, f"_8^: {e}")
    
    async def test_send_messages(self):
        """KmՋ3��S
��mo`"""
        logger.info("KmՋ3��S
��mo`")
        
        try:
            data_message = {
                "command": "test",
                "data": {"message": "Hello WebSocket"}
            }
            
            success = await websocket_manager.send_message(
                client_name=self.test_client_name,
                message_type="data",
                data=data_message,
                priority=0
            )
            
            if success:
                self.add_test_result("�S
�penc�mo`", True, "penc�mo`�S
�b�R")
            else:
                self.add_test_result("�S
�penc�mo`", False, "penc�mo`�S
�1Y%�")
            
            command_message = {
                "action": "ping",
                "timestamp": datetime.now().isoformat()
            }
            
            success = await websocket_manager.send_message(
                client_name=self.test_client_name,
                message_type="command",
                data=command_message,
                priority=1
            )
            
            if success:
                self.add_test_result("�S
�}T�N�mo`", True, "}T�N�mo`�S
�b�R")
            else:
                self.add_test_result("�S
�}T�N�mo`", False, "}T�N�mo`�S
�1Y%�")
                
        except Exception as e:
            self.add_test_result("�S
��mo`", False, f"_8^: {e}")
    
    async def test_receive_message_processing(self):
        """KmՋ4��c6e�mo`Yt"""
        logger.info("KmՋ4��c6e�mo`Yt")
        
        try:
            receive_channel = websocket_manager.get_channel(f"{self.test_client_name}_receive")
            
            if receive_channel:
                test_message = ChannelMessage(
                    type="test_response",
                    data={
                        "status": "success",
                        "message": "Test response received"
                    },
                    priority=0
                )
                
                success = await receive_channel.send_message(test_message)
                
                if success:
                    self.add_test_result("�c6e�mo`Yt", True, "KmՋ�mo`�]�m�R0R�c6eChannel")
                    
                    await asyncio.sleep(0.5)
                    
                    if receive_channel.queue_size == 0:
                        self.add_test_result("�mo`Yt", True, "�mo`�]���c6eYthVYt")
                    else:
                        self.add_test_result("�mo`Yt", False, f"�mo`*g��Yt��R'Y\: {receive_channel.queue_size}")
                else:
                    self.add_test_result("�c6e�mo`Yt", False, "�e�l�m�RKmՋ�mo`0R�c6eChannel")
            else:
                self.add_test_result("�c6e�mo`Yt", False, "�c6eChannel
NX[(W")
                
        except Exception as e:
            self.add_test_result("�c6e�mo`Yt", False, f"_8^: {e}")
    
    async def test_heartbeat_functionality(self):
        """KmՋ5��_���R��"""
        logger.info("KmՋ5��_���R��")
        
        try:
            is_running = websocket_manager._heartbeat_manager.is_heartbeat_running(self.test_client_name)
            
            if is_running:
                self.add_test_result("�_���N�R�r
`", True, "�_���N�Rck(WЏL�")
            else:
                self.add_test_result("�_���N�R�r
`", False, "�_���N�R*gЏL�")
            
            success = await websocket_manager.send_heartbeat(self.test_client_name)
            
            if success:
                self.add_test_result("Kb�R�S
��_��", True, "�_���S
�b�R")
            else:
                self.add_test_result("Kb�R�S
��_��", False, "�_���S
�1Y%�")
                
        except Exception as e:
            self.add_test_result("�_���R��", False, f"_8^: {e}")
    
    async def test_priority_control(self):
        """KmՋ6�OHQ�~�c6R"""
        logger.info("KmՋ6�OHQ�~�c6R")
        
        try:
            send_channel = websocket_manager.get_channel(f"{self.test_client_name}_send")
            heartbeat_channel = websocket_manager.get_channel(f"{self.test_client_name}_heartbeat")
            
            if send_channel and heartbeat_channel:
                high_priority_message = ChannelMessage(
                    type="high_priority",
                    data={"priority": "high"},
                    priority=1
                )
                
                success = await send_channel.send_message(high_priority_message)
                
                if success:
                    self.add_test_result("ؚOHQ�~�mo`", True, "ؚOHQ�~�mo`�m�Rb�R")
                    
                    if send_channel.priority_queue_size > 0:
                        self.add_test_result("OHQ�~�R", True, "�mo`cknxۏeQOHQ�~�R")
                    else:
                        self.add_test_result("OHQ�~�R", False, "�mo`*gۏeQOHQ�~�R")
                else:
                    self.add_test_result("ؚOHQ�~�mo`", False, "ؚOHQ�~�mo`�m�R1Y%�")
            else:
                self.add_test_result("OHQ�~�c6R", False, "�S
�Channelb�_��Channel
NX[(W")
                
        except Exception as e:
            self.add_test_result("OHQ�~�c6R", False, f"_8^: {e}")
    
    async def test_disconnect_client(self):
        """KmՋ7��e_ޏ�c"""
        logger.info("KmՋ7��e_ޏ�c")
        
        try:
            success = await websocket_manager.disconnect_client(self.test_client_name)
            
            if success:
                self.add_test_result("�e_ޏ�c", True, "�[7b�z�e_b�R")
                
                client = websocket_manager.get_client(self.test_client_name)
                if client and not client.is_connected:
                    self.add_test_result("�e_�r
`", True, "�[7b�z�]�e_ޏ�c")
                else:
                    self.add_test_result("�e_�r
`", False, "�[7b�z�NY�Nޏ�c�r
`")
            else:
                self.add_test_result("�e_ޏ�c", False, "�[7b�z�e_1Y%�")
                
        except Exception as e:
            self.add_test_result("�e_ޏ�c", False, f"_8^: {e}")
    
    async def test_cleanup(self):
        """KmՋ8�ntD��n"""
        logger.info("KmՋ8�ntD��n")
        
        try:
            success = await websocket_manager.remove_client(self.test_client_name)
            
            if success:
                self.add_test_result("�yd��[7b�z", True, "�[7b�z�yd�b�R")
                
                client = websocket_manager.get_client(self.test_client_name)
                if client is None:
                    self.add_test_result("�[7b�znt", True, "�[7b�z�]�[hQ�yd�")
                else:
                    self.add_test_result("�[7b�znt", False, "�[7b�z�NX[(W")
            else:
                self.add_test_result("�yd��[7b�z", False, "�[7b�z�yd�1Y%�")
                
        except Exception as e:
            self.add_test_result("ntD��n", False, f"_8^: {e}")
    
    def add_test_result(self, test_name: str, success: bool, message: str):
        """�m�RKmՋ�~�g"""
        result = {
            "test_name": test_name,
            "success": success,
            "message": message,
            "timestamp": datetime.now().isoformat()
        }
        self.test_results.append(result)
        
        status = "' �Ǐ" if success else "L' 1Y%�"
        logger.info(f"{status} {test_name}: {message}")
    
    def print_test_results(self):
        """SbpSKmՋ�~�g"""
        print("\n" + "="*60)
        print("WebSocket�[te�R��KmՋ�~�g")
        print("="*60)
        
        passed = 0
        failed = 0
        
        for result in self.test_results:
            status = "' �Ǐ" if result["success"] else "L' 1Y%�"
            print(f"{status} {result['test_name']}: {result['message']}")
            
            if result["success"]:
                passed += 1
            else:
                failed += 1
        
        print("\n" + "-"*60)
        print(f";`��: {len(self.test_results)} *NKmՋ")
        print(f"�Ǐ: {passed} *N")
        print(f"1Y%�: {failed} *N")
        print(f"b�R�s: {passed/len(self.test_results)*100:.1f}%" if self.test_results else "0%")
        
        if failed == 0:
            print("\n<؉� @b	gKmՋ�Ǐ
�WebSocket�g�g�]\Ock8^0")
        else:
            print(f"\n�&�  	g {failed} *NKmՋ1Y%������h�g�vsQ�R��0")
        
        print("="*60)


async def main():
    """;N�Qpe"""
    print("WebSocket�[te�R��KmՋ
g�R")
    print("="*60)
    
    test = WebSocketCompleteTest()
    await test.run_all_tests()


if __name__ == "__main__":
    asyncio.run(main())