Misc test harnesses (#749)

Some misc test harnesses for a few features

dev-tools/tests/relay/test_rev_gateway.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+"""
+WebSocket Test Client
+
+A simple client to test the reverse gateway through the relay.
+Connects to the relay's /in endpoint and allows sending test messages.
+
+Usage:
+    python test_client.py [--uri URI] [--interactive]
+"""
+
+import asyncio
+import json
+import logging
+import argparse
+import uuid
+from aiohttp import ClientSession, WSMsgType
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("test_client")
+
+class TestClient:
+    """Simple WebSocket test client"""
+    
+    def __init__(self, uri: str):
+        self.uri = uri
+        self.session = None
+        self.ws = None
+        self.running = False
+        self.message_counter = 0
+        self.client_id = str(uuid.uuid4())[:8]
+        
+    async def connect(self):
+        """Connect to the WebSocket"""
+        self.session = ClientSession()
+        logger.info(f"Connecting to {self.uri}")
+        self.ws = await self.session.ws_connect(self.uri)
+        logger.info("Connected successfully")
+        
+    async def disconnect(self):
+        """Disconnect from WebSocket"""
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+        if self.session and not self.session.closed:
+            await self.session.close()
+        logger.info("Disconnected")
+        
+    async def send_message(self, service: str, request_data: dict, flow: str = "default"):
+        """Send a properly formatted TrustGraph message"""
+        self.message_counter += 1
+        message = {
+            "id": f"{self.client_id}-{self.message_counter}",
+            "service": service,
+            "request": request_data,
+            "flow": flow
+        }
+        
+        message_json = json.dumps(message, indent=2)
+        logger.info(f"Sending message:\n{message_json}")
+        await self.ws.send_str(json.dumps(message))
+        
+    async def listen_for_responses(self):
+        """Listen for incoming messages"""
+        logger.info("Listening for responses...")
+        
+        async for msg in self.ws:
+            if msg.type == WSMsgType.TEXT:
+                try:
+                    response = json.loads(msg.data)
+                    logger.info(f"Received response:\n{json.dumps(response, indent=2)}")
+                except json.JSONDecodeError:
+                    logger.info(f"Received text: {msg.data}")
+            elif msg.type == WSMsgType.BINARY:
+                logger.info(f"Received binary data: {len(msg.data)} bytes")
+            elif msg.type == WSMsgType.ERROR:
+                logger.error(f"WebSocket error: {self.ws.exception()}")
+                break
+            else:
+                logger.info(f"Connection closed: {msg.type}")
+                break
+                
+    async def interactive_mode(self):
+        """Interactive mode for manual testing"""
+        print("\n=== Interactive Test Client ===")
+        print("Available commands:")
+        print("  text-completion - Test text completion service")
+        print("  agent          - Test agent service")
+        print("  embeddings     - Test embeddings service")
+        print("  custom         - Send custom message")
+        print("  quit           - Exit")
+        print()
+        
+        # Start response listener
+        listen_task = asyncio.create_task(self.listen_for_responses())
+        
+        try:
+            while True:
+                try:
+                    command = input("Command> ").strip().lower()
+                    
+                    if command == "quit":
+                        break
+                    elif command == "text-completion":
+                        await self.send_message("text-completion", {
+                            "system": "You are a helpful assistant.",
+                            "prompt": "What is 2+2?"
+                        })
+                    elif command == "agent":
+                        await self.send_message("agent", {
+                            "question": "What is the capital of France?"
+                        })
+                    elif command == "embeddings":
+                        await self.send_message("embeddings", {
+                            "text": "Hello world"
+                        })
+                    elif command == "custom":
+                        service = input("Service name> ").strip()
+                        request_json = input("Request JSON> ").strip()
+                        try:
+                            request_data = json.loads(request_json)
+                            await self.send_message(service, request_data)
+                        except json.JSONDecodeError as e:
+                            print(f"Invalid JSON: {e}")
+                    elif command == "":
+                        continue
+                    else:
+                        print(f"Unknown command: {command}")
+                        
+                except KeyboardInterrupt:
+                    break
+                except EOFError:
+                    break
+                except Exception as e:
+                    logger.error(f"Error in interactive mode: {e}")
+                    
+        finally:
+            listen_task.cancel()
+            try:
+                await listen_task
+            except asyncio.CancelledError:
+                pass
+            
+    async def run_predefined_tests(self):
+        """Run a series of predefined tests"""
+        print("\n=== Running Predefined Tests ===")
+        
+        # Start response listener
+        listen_task = asyncio.create_task(self.listen_for_responses())
+        
+        try:
+            # Test 1: Text completion
+            print("\n1. Testing text-completion service...")
+            await self.send_message("text-completion", {
+                "system": "You are a helpful assistant.",
+                "prompt": "What is 2+2?"
+            })
+            await asyncio.sleep(2)
+            
+            # Test 2: Agent
+            print("\n2. Testing agent service...")
+            await self.send_message("agent", {
+                "question": "What is the capital of France?"
+            })
+            await asyncio.sleep(2)
+            
+            # Test 3: Embeddings
+            print("\n3. Testing embeddings service...")
+            await self.send_message("embeddings", {
+                "text": "Hello world"
+            })
+            await asyncio.sleep(2)
+            
+            # Test 4: Invalid service
+            print("\n4. Testing invalid service...")
+            await self.send_message("nonexistent-service", {
+                "test": "data"
+            })
+            await asyncio.sleep(2)
+            
+            print("\nTests completed. Waiting for any remaining responses...")
+            await asyncio.sleep(3)
+            
+        finally:
+            listen_task.cancel()
+            try:
+                await listen_task
+            except asyncio.CancelledError:
+                pass
+
+async def main():
+    parser = argparse.ArgumentParser(
+        description="WebSocket Test Client for Reverse Gateway"
+    )
+    parser.add_argument(
+        '--uri',
+        default='ws://localhost:8080/in',
+        help='WebSocket URI to connect to (default: ws://localhost:8080/in)'
+    )
+    parser.add_argument(
+        '--interactive', '-i',
+        action='store_true',
+        help='Run in interactive mode'
+    )
+    parser.add_argument(
+        '--verbose', '-v',
+        action='store_true',
+        help='Enable verbose logging'
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    client = TestClient(args.uri)
+    
+    try:
+        await client.connect()
+        
+        if args.interactive:
+            await client.interactive_mode()
+        else:
+            await client.run_predefined_tests()
+            
+    except KeyboardInterrupt:
+        print("\nShutdown requested by user")
+    except Exception as e:
+        logger.error(f"Client error: {e}")
+    finally:
+        await client.disconnect()
+
+if __name__ == "__main__":
+    asyncio.run(main())

dev-tools/tests/relay/websocket_relay.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python3
+"""
+WebSocket Relay Test Harness
+
+This script creates a relay server with two WebSocket endpoints:
+- /in  - for test clients to connect to
+- /out - for reverse gateway to connect to
+
+Messages are bidirectionally relayed between the two connections.
+
+Usage:
+    python websocket_relay.py [--port PORT] [--host HOST]
+"""
+
+import asyncio
+import logging
+import argparse
+from aiohttp import web, WSMsgType
+import weakref
+from typing import Optional, Set
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("websocket_relay")
+
+class WebSocketRelay:
+    """WebSocket relay that forwards messages between 'in' and 'out' connections"""
+    
+    def __init__(self):
+        self.in_connections: Set = weakref.WeakSet()
+        self.out_connections: Set = weakref.WeakSet()
+        
+    async def handle_in_connection(self, request):
+        """Handle incoming connections on /in endpoint"""
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        
+        self.in_connections.add(ws)
+        logger.info(f"New 'in' connection. Total in: {len(self.in_connections)}, out: {len(self.out_connections)}")
+        
+        try:
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    data = msg.data
+                    logger.info(f"IN → OUT: {data}")
+                    await self._forward_to_out(data)
+                elif msg.type == WSMsgType.BINARY:
+                    data = msg.data
+                    logger.info(f"IN → OUT: {len(data)} bytes (binary)")
+                    await self._forward_to_out(data, binary=True)
+                elif msg.type == WSMsgType.ERROR:
+                    logger.error(f"WebSocket error on 'in' connection: {ws.exception()}")
+                    break
+                else:
+                    break
+        except Exception as e:
+            logger.error(f"Error in 'in' connection handler: {e}")
+        finally:
+            logger.info(f"'in' connection closed. Remaining in: {len(self.in_connections)}, out: {len(self.out_connections)}")
+            
+        return ws
+    
+    async def handle_out_connection(self, request):
+        """Handle outgoing connections on /out endpoint"""
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        
+        self.out_connections.add(ws)
+        logger.info(f"New 'out' connection. Total in: {len(self.in_connections)}, out: {len(self.out_connections)}")
+        
+        try:
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    data = msg.data
+                    logger.info(f"OUT → IN: {data}")
+                    await self._forward_to_in(data)
+                elif msg.type == WSMsgType.BINARY:
+                    data = msg.data
+                    logger.info(f"OUT → IN: {len(data)} bytes (binary)")
+                    await self._forward_to_in(data, binary=True)
+                elif msg.type == WSMsgType.ERROR:
+                    logger.error(f"WebSocket error on 'out' connection: {ws.exception()}")
+                    break
+                else:
+                    break
+        except Exception as e:
+            logger.error(f"Error in 'out' connection handler: {e}")
+        finally:
+            logger.info(f"'out' connection closed. Remaining in: {len(self.in_connections)}, out: {len(self.out_connections)}")
+            
+        return ws
+    
+    async def _forward_to_out(self, data, binary=False):
+        """Forward message from 'in' to all 'out' connections"""
+        if not self.out_connections:
+            logger.warning("No 'out' connections available to forward message")
+            return
+            
+        closed_connections = []
+        for ws in list(self.out_connections):
+            try:
+                if ws.closed:
+                    closed_connections.append(ws)
+                    continue
+                    
+                if binary:
+                    await ws.send_bytes(data)
+                else:
+                    await ws.send_str(data)
+            except Exception as e:
+                logger.error(f"Error forwarding to 'out' connection: {e}")
+                closed_connections.append(ws)
+        
+        # Clean up closed connections
+        for ws in closed_connections:
+            if ws in self.out_connections:
+                self.out_connections.discard(ws)
+    
+    async def _forward_to_in(self, data, binary=False):
+        """Forward message from 'out' to all 'in' connections"""
+        if not self.in_connections:
+            logger.warning("No 'in' connections available to forward message")
+            return
+            
+        closed_connections = []
+        for ws in list(self.in_connections):
+            try:
+                if ws.closed:
+                    closed_connections.append(ws)
+                    continue
+                    
+                if binary:
+                    await ws.send_bytes(data)
+                else:
+                    await ws.send_str(data)
+            except Exception as e:
+                logger.error(f"Error forwarding to 'in' connection: {e}")
+                closed_connections.append(ws)
+        
+        # Clean up closed connections
+        for ws in closed_connections:
+            if ws in self.in_connections:
+                self.in_connections.discard(ws)
+
+async def create_app(relay):
+    """Create the web application with routes"""
+    app = web.Application()
+    
+    # Add routes
+    app.router.add_get('/in', relay.handle_in_connection)
+    app.router.add_get('/out', relay.handle_out_connection)
+    
+    # Add a simple status endpoint
+    async def status(request):
+        status_info = {
+            'in_connections': len(relay.in_connections),
+            'out_connections': len(relay.out_connections),
+            'status': 'running'
+        }
+        return web.json_response(status_info)
+    
+    app.router.add_get('/status', status)
+    app.router.add_get('/', status)  # Root also shows status
+    
+    return app
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="WebSocket Relay Test Harness"
+    )
+    parser.add_argument(
+        '--host', 
+        default='localhost',
+        help='Host to bind to (default: localhost)'
+    )
+    parser.add_argument(
+        '--port', 
+        type=int, 
+        default=8080,
+        help='Port to bind to (default: 8080)'
+    )
+    parser.add_argument(
+        '--verbose', '-v',
+        action='store_true',
+        help='Enable verbose logging'
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    relay = WebSocketRelay()
+    
+    print(f"Starting WebSocket Relay on {args.host}:{args.port}")
+    print(f"  'in' endpoint:  ws://{args.host}:{args.port}/in")
+    print(f"  'out' endpoint: ws://{args.host}:{args.port}/out")
+    print(f"  Status:         http://{args.host}:{args.port}/status")
+    print()
+    print("Usage:")
+    print(f"  Test client connects to:     ws://{args.host}:{args.port}/in")
+    print(f"  Reverse gateway connects to: ws://{args.host}:{args.port}/out")
+    
+    web.run_app(create_app(relay), host=args.host, port=args.port)
+
+if __name__ == "__main__":
+    main()