Improved mock DevRev Server

jest.setup.js

@@ -76,11 +76,59 @@ global.test = (name, fn, timeout) => {
   }, timeout);
 };
 
+// Function to reset the mock server state
+const resetMockServer = async () => {
+  try {
+    const http = require('http');
+
+    // Simple HTTP request that properly closes the connection
+    const postData = '';
+    const options = {
+      hostname: 'localhost',
+      port: 8003,
+      path: '/reset-mock-server',
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Content-Length': Buffer.byteLength(postData)
+      }
+    };
+
+    const response = await new Promise((resolve, reject) => {
+      const req = http.request(options, (res) => {
+        let data = '';
+        res.on('data', (chunk) => data += chunk);
+        res.on('end', () => resolve({ status: res.statusCode, data }));
+      });
+
+      req.on('error', reject);
+      req.on('timeout', () => {
+        req.destroy();
+        reject(new Error('Request timeout'));
+      });
+
+      req.setTimeout(5000);
+      req.write(postData);
+      req.end();
+    });
+
+    if (response.status !== 200) {
+      console.warn(`Failed to reset mock server: ${response.status}`);
+    }
+  } catch (error) {
+    console.warn(`Could not connect to mock server for reset: ${error.message}`);
+  }
+};
+
 const originalBeforeEach = global.beforeEach;
 
 global.beforeEach = (fn, timeout) => {
   originalBeforeEach(async () => {
     try {
+      // Reset the mock server state before each test
+      await resetMockServer();
+
+      // Execute the original beforeEach function
       await fn();
     } catch (err) {
       expect(`beforeEach failed: ${err.message}`).toBe('beforeEach should not fail');

mock_devrev_server.py

@@ -4,9 +4,14 @@
 import uuid
 import json
 import random
+import copy
 
 app = FastAPI()
 
+# Initialize application state containers
+app.state.uploaded_states = {}
+app.state.uploaded_artifacts = set()
+
 class ArtifactPrepareRequest(BaseModel):
     file_name: str
     file_type: str
@@ -28,6 +33,13 @@ class AirdropArtifactResponse(BaseModel):
     upload_url: str
     form_data: Dict[str, str]
 
+@app.get("/is_uploaded/{artifact_id:path}")
+async def was_artifact_uploaded(artifact_id: str):
+    """Check if an artifact with the given artifact_id was uploaded"""
+    if artifact_id in app.state.uploaded_artifacts:
+        return {"artifact_id": artifact_id, "uploaded": True}
+    raise HTTPException(status_code=404, detail="Artifact not found")
+
 @app.post("/upload/{artifact_id:path}")
 async def upload_artifact(
     artifact_id: str,
@@ -41,6 +53,9 @@ async def upload_artifact(
         print(f"Content length: {len(content)}")
         print(f"Content type: {request.headers.get('content-type', 'unknown')}")
 
+        # Remember that this artifact_id was uploaded
+        app.state.uploaded_artifacts.add(artifact_id)
+
         return {"status": "success", "message": "File uploaded successfully"}
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
@@ -70,9 +85,11 @@ async def prepare_artifact(
     )
 
 @app.get("/external-worker.get", response_model=ExternalWorkerResponse)
-async def get_external_worker():
-    print("Received /external-worker.get GET body:")
-    state = {
+async def get_external_worker(sync_unit: str):
+    print(f"Received /external-worker.get GET request for sync_unit: {sync_unit}")
+
+    # Default state
+    default_state = {
         "lastSyncStarted": "2024-06-01T12:00:00Z",
         "toDevRev": {
             "attachmentsMetadata": {
@@ -89,14 +106,24 @@ async def get_external_worker():
             "offset": 0
         }
     }
-    return ExternalWorkerResponse(state=json.dumps(state))
+
+    # Check if uploaded_states contains the specific sync_unit
+    if sync_unit in app.state.uploaded_states:
+        print(f"Found uploaded state for sync_unit: {sync_unit}")
+        return ExternalWorkerResponse(state=json.dumps(app.state.uploaded_states[sync_unit]))
+    else:
+        print(f"No uploaded state found for sync_unit: {sync_unit}, returning default state")
+        return ExternalWorkerResponse(state=json.dumps(default_state))
 
 @app.post("/external-worker.update")
-async def update_external_worker(request: Request):
+async def update_external_worker(sync_unit: str, request: Request):
     body = await request.body()
-    print("Received /external-worker.update POST body:")
+    print(f"Received /external-worker.update POST request for sync_unit: {sync_unit}")
     try:
         parsed = json.loads(body.decode("utf-8"))
+        # Store the uploaded state under the specific sync_unit key
+        app.state.uploaded_states[sync_unit] = copy.deepcopy(parsed)
+        print(f"Stored state for sync_unit: {sync_unit}")
         print(json.dumps(parsed, indent=2))
     except Exception as e:
         print("Failed to pretty print JSON:", e)
@@ -148,6 +175,14 @@ async def confirm_upload(request: Request):
 
     return {"status": "success"}
 
+@app.post("/reset-mock-server")
+async def reset_mock_server():
+    """Reset the mock server state by clearing uploaded_states and uploaded_artifacts"""
+    app.state.uploaded_states = {}
+    app.state.uploaded_artifacts = set()
+    print("Mock server state reset - uploaded_states and uploaded_artifacts cleared")
+    return {"status": "success", "message": "Mock server state reset successfully"}
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="localhost", port=8003)