From e2e87ee691d73ea5063b0f9dcbbd4f6d61e85b34 Mon Sep 17 00:00:00 2001 From: "agentfarmx[bot]" <198411105+agentfarmx[bot]@users.noreply.github.com> Date: Sat, 1 Mar 2025 08:59:31 +0000 Subject: [PATCH] feat: enhance test_tcp.py with rich text formatting for improved output --- test_tcp.py | 69 +++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 49 insertions(+), 20 deletions(-) diff --git a/test_tcp.py b/test_tcp.py index 505fd3e..51c3642 100755 --- a/test_tcp.py +++ b/test_tcp.py @@ -2,6 +2,14 @@ import socket import json import sys +from rich.console import Console +from rich.panel import Panel +from rich.syntax import Syntax +from rich.table import Table +from rich import print as rprint + +# Create a Rich console object +console = Console() def send_receive(sock, request_dict): """Send a request to the server and receive the response.""" @@ -60,6 +68,26 @@ def receive_response(sock): response_json = b''.join(chunks).decode('utf-8') return json.loads(response_json) +def display_response(response, title="Response"): + """Display a response using Rich formatting.""" + # Create a table for the response + table = Table(title=title) + table.add_column("Key", style="cyan") + table.add_column("Value", style="green") + + for key, value in response.items(): + if key == "output": + # Format output as a panel + table.add_row(key, Panel(value.strip(), border_style="blue")) + elif key == "error": + # Format error with red text + table.add_row(key, f"[bold red]{value}[/bold red]") + else: + # Format other values normally + table.add_row(key, str(value)) + + console.print(table) + def main(): # Server connection details host = "localhost" @@ -69,33 +97,34 @@ def main(): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) try: sock.connect((host, port)) - print(f"Connected to {host}:{port}") + console.print(f"Connected to [bold]{host}:{port}[/bold]", style="green") # Receive initial greeting using the length-prefixed protocol greeting = receive_response(sock) - print(f"Server greeting: {json.dumps(greeting, indent=2)}") + console.print(Panel(Syntax(json.dumps(greeting, indent=2), "json", theme="monokai"), + title="Server Greeting", border_style="green")) # Test 1: Execute code without a session ID (creates a new session) - print("\n--- Test 1: Execute code without a session ID ---") + console.rule("[bold]Test 1: Execute code without a session ID[/bold]") response = send_receive(sock, { "code": "x = 42\nprint(f'x = {x}')" }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 1 Response") # Save the session ID for later use session_id = response.get("session_id") - print(f"Session ID: {session_id}") + console.print(f"Session ID: [bold cyan]{session_id}[/bold cyan]") # Test 2: Execute code in the same session (using the session ID) - print("\n--- Test 2: Execute code in the same session ---") + console.rule("[bold]Test 2: Execute code in the same session[/bold]") response = send_receive(sock, { "code": "y = x * 2\nprint(f'y = {y}')", "session_id": session_id }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 2 Response") # Test 3: Define a function in the session - print("\n--- Test 3: Define a function in the session ---") + console.rule("[bold]Test 3: Define a function in the session[/bold]") response = send_receive(sock, { "code": """ def greet(name): @@ -104,46 +133,46 @@ def greet(name): """, "session_id": session_id }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 3 Response") # Test 4: Call the function defined in the previous request - print("\n--- Test 4: Call the function defined in the previous request ---") + console.rule("[bold]Test 4: Call the function defined in the previous request[/bold]") response = send_receive(sock, { "code": "print(greet('Python'))", "session_id": session_id }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 4 Response") # Test 5: Create a new session - print("\n--- Test 5: Create a new session ---") + console.rule("[bold]Test 5: Create a new session[/bold]") response = send_receive(sock, { "code": "print('This is a new session')" }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 5 Response") new_session_id = response.get("session_id") - print(f"New Session ID: {new_session_id}") + console.print(f"New Session ID: [bold cyan]{new_session_id}[/bold cyan]") # Test 6: Verify the new session doesn't have access to variables from the first session - print("\n--- Test 6: Verify session isolation ---") + console.rule("[bold]Test 6: Verify session isolation[/bold]") response = send_receive(sock, { "code": "try:\n print(f'x = {x}')\nexcept NameError as e:\n print(f'Error: {e}')", "session_id": new_session_id }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 6 Response") # Test 7: Try to access a non-existent session - print("\n--- Test 7: Try to access a non-existent session ---") + console.rule("[bold]Test 7: Try to access a non-existent session[/bold]") response = send_receive(sock, { "code": "print('This should fail')", "session_id": "non-existent-session-id" }) - print(f"Response: {json.dumps(response, indent=2)}") + display_response(response, "Test 7 Response") except Exception as e: - print(f"Error: {e}") + console.print(f"[bold red]Error:[/bold red] {e}") finally: sock.close() - print("Connection closed") + console.print("Connection closed", style="yellow") if __name__ == "__main__": main() \ No newline at end of file