Writing tests to break the server

server.py

@@ -2,8 +2,9 @@
 import json
 import numpy as np
 import uvicorn
+from pydantic import BaseModel
 from pydantic_settings import BaseSettings
-from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect
+from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect, HTTPException
 from datetime import datetime
 import msgpack
 import asyncio
@@ -12,9 +13,14 @@
 from contextlib import asynccontextmanager
 
 
+class CloseRequest(BaseModel):
+    reason: Optional[str] = None
+
+
 class Settings(BaseSettings):
     redis_url: str = "redis://localhost:6379/0"
     ttl: int = 60 * 60  # 1 hour
+    max_payload_size: int = 16 * 1024 * 1024  # 16MB max payload
 
 
 def build_app(settings: Settings):
@@ -57,9 +63,16 @@ async def close(node_id):
     async def append(node_id, request: Request):
         "Append data to a dataset."
 
+        # Check request body size limit
+        # Tell good-faith clients that their request is too big.
+        # Fix for: test_large_data_resource.py::test_large_data_resource_limits
+        headers = request.headers
+        content_length = headers.get("content-length")
+        if content_length and int(content_length) > settings.max_payload_size:
+            raise HTTPException(status_code=413, detail="Payload too large")
+
         # get data from request body
         binary_data = await request.body()
-        headers = request.headers
         metadata = {
             "timestamp": datetime.now().isoformat(),
         }
@@ -86,12 +99,11 @@ async def append(node_id, request: Request):
     #   @app.websocket("/stream/many")
 
     @app.post("/close/{node_id}")
-    async def close_connection(node_id: str, request: Request):
-        # Parse the JSON body
-        body = await request.json()
+    async def close_connection(node_id: str, body: CloseRequest, request: Request):
+        # Pydantic automatically handles JSON parsing and validation
+        # Fix for: test_json_parsing.py::test_json_parsing_errors_in_close_endpoint
         headers = request.headers
-
-        reason = body.get("reason", None)
+        reason = body.reason
 
         metadata = {"timestamp": datetime.now().isoformat(), "reason": reason}
         metadata.setdefault("Content-Type", headers.get("Content-Type"))

tests/conftest.py

@@ -8,7 +8,6 @@ def client():
     """Fixture providing TestClient following ws-tests pattern."""
     settings = Settings(redis_url="redis://localhost:6379/0", ttl=60 * 60)
     app = build_app(settings)
-    
+
     with TestClient(app) as client:
         yield client
-

tests/test_json_parsing.py

@@ -0,0 +1,24 @@
+"""
+Tests for JSON parsing error handling bugs in server endpoints.
+"""
+
+
+def test_json_parsing_errors_in_close_endpoint(client):
+    """Server should handle JSON parsing errors in /close endpoint gracefully."""
+    response = client.post("/upload")
+    assert response.status_code == 200
+    node_id = response.json()["node_id"]
+
+    # Test 1: Malformed JSON content should not crash
+    response = client.post(
+        f"/close/{node_id}",
+        content=b"invalid json {{{",
+        headers={"Content-Type": "application/json"},
+    )
+    assert response.status_code == 422  # Pydantic returns 422 for validation errors
+    assert "detail" in response.json()
+
+    # Test 2: Missing JSON body should not crash
+    response = client.post(f"/close/{node_id}")
+    assert response.status_code == 422  # Pydantic returns 422 for validation errors
+    assert "detail" in response.json()

tests/test_large_data_resource.py

@@ -0,0 +1,22 @@
+"""
+Tests for large data handling and resource limit bugs.
+"""
+
+
+def test_large_data_resource_limits(client):
+    """Server should handle large data with proper resource limits."""
+
+    # Test: Huge payload (20MB) - should be rejected as too large
+    response = client.post("/upload")
+    assert response.status_code == 200
+    node_id = response.json()["node_id"]
+
+    huge_payload = b"\x00" * (20 * 1024 * 1024)  # 20MB (exceeds 16MB limit)
+    response = client.post(
+        f"/upload/{node_id}",
+        content=huge_payload,
+        headers={"Content-Type": "application/octet-stream"},
+    )
+    # Should be rejected with 413 Payload Too Large due to size limits
+    assert response.status_code == 413
+    assert "Payload too large" in response.json()["detail"]