city/tests/test_client.py

"""
WebSocket test client for simulating real game interactions
"""
import asyncio
import json
import websockets
from typing import Dict, List, Optional, Any
from contextlib import asynccontextmanager


class TestWebSocketClient:
    """Test client that simulates a real game client WebSocket connection"""
    
    def __init__(self, nickname: str, base_url: str = "ws://127.0.0.1:9901"):
        self.nickname = nickname
        self.base_url = base_url
        self.websocket = None
        self.received_messages = []
        self.player_data = None
        self.game_state = None
        self.connected = False
        
    async def connect(self):
        """Connect to the WebSocket server"""
        try:
            uri = f"{self.base_url}/ws/{self.nickname}"
            self.websocket = await websockets.connect(uri)
            self.connected = True
            
            # Wait for initial game state
            init_message = await self.websocket.recv()
            init_data = json.loads(init_message)
            
            if init_data["type"] == "init":
                self.player_data = init_data["player"]
                self.game_state = init_data["game_state"]
                
            return True
        except Exception as e:
            print(f"Failed to connect {self.nickname}: {e}")
            return False
    
    async def disconnect(self):
        """Disconnect from the WebSocket server"""
        if self.websocket:
            await self.websocket.close()
            self.connected = False
    
    async def send_message(self, message_type: str, data: Dict[str, Any]):
        """Send a message to the server"""
        if not self.connected:
            raise ConnectionError("Client not connected")
        
        message = {"type": message_type, **data}
        await self.websocket.send(json.dumps(message))
    
    async def receive_message(self, timeout: float = 1.0) -> Optional[Dict]:
        """Receive a message from the server with timeout"""
        if not self.connected:
            return None
            
        try:
            message = await asyncio.wait_for(self.websocket.recv(), timeout=timeout)
            data = json.loads(message)
            self.received_messages.append(data)
            return data
        except asyncio.TimeoutError:
            return None
    
    async def receive_messages_for(self, duration: float) -> List[Dict]:
        """Collect all messages received during a time period"""
        messages = []
        start_time = asyncio.get_event_loop().time()
        
        while (asyncio.get_event_loop().time() - start_time) < duration:
            try:
                remaining_time = duration - (asyncio.get_event_loop().time() - start_time)
                if remaining_time <= 0:
                    break
                    
                message = await asyncio.wait_for(self.websocket.recv(), timeout=remaining_time)
                data = json.loads(message)
                messages.append(data)
                self.received_messages.append(data)
            except asyncio.TimeoutError:
                break
        
        return messages
    
    # Game action methods
    async def place_building(self, building_type: str, x: int, y: int):
        """Place a building at coordinates"""
        await self.send_message("place_building", {
            "building_type": building_type,
            "x": x,
            "y": y
        })
    
    async def remove_building(self, x: int, y: int):
        """Remove a building at coordinates"""
        await self.send_message("remove_building", {"x": x, "y": y})
    
    async def edit_building(self, x: int, y: int, name: str):
        """Edit a building name"""
        await self.send_message("edit_building", {"x": x, "y": y, "name": name})
    
    async def send_chat(self, message: str):
        """Send a chat message"""
        await self.send_message("chat", {
            "message": message,
            "timestamp": "12:00"
        })
    
    async def send_cursor_move(self, x: int, y: int):
        """Send cursor movement"""
        await self.send_message("cursor_move", {"x": x, "y": y})
    
    def get_money(self) -> int:
        """Get current player money"""
        return self.player_data["money"] if self.player_data else 0
    
    def get_population(self) -> int:
        """Get current player population"""
        return self.player_data["population"] if self.player_data else 0
    
    def clear_messages(self):
        """Clear received message history"""
        self.received_messages.clear()


@asynccontextmanager
async def test_clients(*nicknames):
    """Context manager to create and manage multiple test clients"""
    clients = []
    
    try:
        # Create and connect all clients
        for nickname in nicknames:
            client = TestWebSocketClient(nickname)
            if await client.connect():
                clients.append(client)
            else:
                raise ConnectionError(f"Failed to connect client {nickname}")
        
        yield clients
        
    finally:
        # Disconnect all clients
        for client in clients:
            if client.connected:
                await client.disconnect()


# Remove test function that was causing warnings
def _test_clients():
    """Placeholder - test_clients is a context manager, not a test"""
    pass


class TestGameServer:
    """Utility class to manage test server lifecycle"""
    
    def __init__(self):
        self.server_process = None
        self.base_url = "http://127.0.0.1:9901"
        self.ws_url = "ws://127.0.0.1:9901"
    
    async def start_server(self):
        """Start the game server for testing"""
        # For testing, we'll assume the server is already running
        # In a more complex setup, we could start/stop the server here
        pass
    
    async def stop_server(self):
        """Stop the game server"""
        pass
    
    async def reset_database(self):
        """Reset the game database for clean tests"""
        import os
        db_path = "/home/retoor/projects/city/data/game.db"
        if os.path.exists(db_path):
            os.remove(db_path)