scripts/test_middleware.py

import asyncio
import json
import os
import sys
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware

# Mock setup
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, os.path.join(project_root, "services", "ai-service", "src"))

from ai_med_extract.app import BenchmarkMiddleware
from ai_med_extract.utils.benchmark import BenchmarkContext

# Mock Request
class MockRequest:
    def __init__(self, method="POST", path="/api/generate"):
        self.method = method
        self.url = type('obj', (object,), {'path': path})

class MockResponse:
    def __init__(self, status_code=200):
        self.status_code = status_code

async def mock_call_next(request):
    await asyncio.sleep(0.1) # Simulate work
    return MockResponse(200)

async def test_middleware():
    mw = BenchmarkMiddleware(None)
    
    print("Testing Middleware...")
    request = MockRequest(method="POST", path="/api/test_generate")
    await mw.dispatch(request, mock_call_next)
    print("Middleware dispatch complete.")

def verify_logs():
    print("Verifying logs...")
    log_file = "logs/benchmark.jsonl"
    if os.path.exists(log_file):
        with open(log_file, "r") as f:
            lines = f.readlines()
            last_line = json.loads(lines[-1])
            print("Last log entry:")
            print(json.dumps(last_line, indent=2))
            
            # Assertions
            if last_line["activity"] == "API_POST_/api/test_generate":
                print("SUCCESS: Activity name matches.")
            else:
                print(f"FAILURE: Activity name mismatch. Got {last_line['activity']}")
            
            if last_line["meta"].get("status_code") == 200:
                 print("SUCCESS: Status code captured.")
            else:
                 print(f"FAILURE: Status code missing or wrong. Got {last_line['meta'].get('status_code')}")

if __name__ == "__main__":
    asyncio.run(test_middleware())
    verify_logs()