import asyncio
import unittest
import websockets
import time

WS_URL = "ws://localhost:8501/_stcore/stream"
CONCURRENT_CLIENTS = 50
TEST_DURATION_SEC = 5

async def simulate_streamlit_client(client_id):
    try:
        async with websockets.connect(WS_URL) as ws:
            msg = await ws.recv()
            start_time = time.time()
            recv_count = 0
            while time.time() - start_time < TEST_DURATION_SEC:
                await ws.ping()
                await asyncio.sleep(0.5)
                recv_count += 1
            return True, recv_count
    except Exception as e:
        return False, str(e)

class TestWebSocketConcurrency(unittest.TestCase):
    def test_ws_concurrency_load(self):
        async def run_test():
            tasks = [simulate_streamlit_client(i) for i in range(CONCURRENT_CLIENTS)]
            return await asyncio.gather(*tasks)

        results = asyncio.run(run_test())
        successes = sum(1 for status, _ in results if status)
        failures = sum(1 for status, _ in results if not status)

        print(f"\n📊 Concurrency Test: {successes}/{CONCURRENT_CLIENTS} Connected Successfully")
        self.assertEqual(failures, 0, f"Failed {failures} WebSocket connections under load.")

if __name__ == "__main__":
    unittest.main()