From 86c5ff2ebc4595fc086b9240101b6fbea088c378 Mon Sep 17 00:00:00 2001
From: Ro'e Katz <roekatz@gmail.com>
Date: Wed, 24 Apr 2024 13:46:03 +0300
Subject: [PATCH] Try fixing multiprocess test

---
 tests/multiprocess_test.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/tests/multiprocess_test.py b/tests/multiprocess_test.py
index 9c5189e..8c4a211 100644
--- a/tests/multiprocess_test.py
+++ b/tests/multiprocess_test.py
@@ -3,6 +3,7 @@
         Publishing-Client -> PubSubServer -> Subscribing->Client
 
 """
+
 import os
 import sys
 import pytest
@@ -35,11 +36,9 @@
 
 def setup_server():
     app = FastAPI()
-    router = APIRouter()
     # PubSub websocket endpoint
     endpoint = PubSubEndpoint()
-    endpoint.register_route(router, path="/pubsub")
-    app.include_router(router)
+    endpoint.register_route(app, path="/pubsub")
     uvicorn.run(app, port=PORT)
 
 
@@ -52,14 +51,14 @@ async def actual():
         # Wait for other client to wake up before publishing to it
         CLIENT_START_SYNC.wait(5)
         # Create a client and subscribe to topics
-        client = PubSubClient()
-        client.start_client(uri)
-        # wait for the client to be ready
-        await client.wait_until_ready()
-        # publish event
-        logger.info("Publishing event")
-        published = await client.publish([EVENT_TOPIC], data=DATA)
-        assert published.result
+        async with PubSubClient() as client:
+            client.start_client(uri)
+            # wait for the client to be ready
+            await client.wait_until_ready()
+            # publish event
+            logger.info("Publishing event")
+            published = await client.publish([EVENT_TOPIC], data=DATA)
+            assert published.result
 
     logger.info("Starting async publishing client")
     asyncio.get_event_loop().run_until_complete(actual())