CoolFace
Apppublic

smart-models/Placebo_AI

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
04_test_load_stress.py43 linesDownload Raw Back to validation_suite
1import asyncio2import time3import requests4import aiohttp5 6BASE_URL = "http://localhost:8000"7# NOTE: You MUST provide a valid JWT token here if your API is secured.8# If testing without Auth, ensure the backend route allows it.9TEST_JWT_TOKEN = "your_test_jwt_token_here" 10 11async def fetch_chat(session, prompt):12    headers = {"Authorization": f"Bearer {TEST_JWT_TOKEN}"}13    payload = {"message": prompt, "role": "MBBS Student/Doctor", "history": []}14    15    start_time = time.time()16    try:17        # Note: Streaming responses make latency calculation complex. 18        # We are just measuring Time-to-First-Token/Connection.19        async with session.post(f"{BASE_URL}/chat", json=payload, headers=headers) as response:20            latency = time.time() - start_time21            return response.status, latency22    except Exception as e:23        return 500, time.time() - start_time24 25async def test_load_stress():26    print("--- Phase 4: Load & Stress Testing ---")27    prompts = ["What is Aspirin?", "Define Appendicitis", "Mechanism of Metformin", "Explain Hypertension"]28    29    print(f"[TEST] Firing {len(prompts)} concurrent requests to stress Test VRAM...")30    async with aiohttp.ClientSession() as session:31        tasks = [fetch_chat(session, p) for p in prompts]32        results = await asyncio.gather(*tasks)33        34        for i, (status, latency) in enumerate(results):35            print(f"Request {i+1}: Status {status} | Latency: {latency:.2f}s")36            37        success_rate = sum(1 for status, _ in results if status == 200) / len(results) * 10038        print(f"--- Load Test Completed ---")39        print(f"Success Rate: {success_rate}%")40        41if __name__ == "__main__":42    asyncio.run(test_load_stress())43