From 14df3a6ffb39e04b2202eb13a0873a8c72d58e3a Mon Sep 17 00:00:00 2001
From: Vlad Stan <stan.v.vlad@gmail.com>
Date: Tue, 7 Feb 2023 18:15:43 +0200
Subject: [PATCH] tests: update after changes

---
 client_manager.py     |  2 +-
 tests/test_clients.py | 24 +++++++++++++++++-------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/client_manager.py b/client_manager.py
index 45f9eb6..6e8f548 100644
--- a/client_manager.py
+++ b/client_manager.py
@@ -62,12 +62,12 @@ class NostrClientManager:
 
 
 class NostrClientConnection:
-    broadcast_event: Callable
 
     def __init__(self, relay_id: str, websocket: WebSocket):
         self.websocket = websocket
         self.relay_id = relay_id
         self.filters: List[NostrFilter] = []
+        self.broadcast_event: Optional[Callable] = None
 
     async def start(self):
         await self.websocket.accept()
diff --git a/tests/test_clients.py b/tests/test_clients.py
index 49fedf7..54a4da2 100644
--- a/tests/test_clients.py
+++ b/tests/test_clients.py
@@ -1,6 +1,7 @@
 import asyncio
-from copy import deepcopy
 from json import dumps, loads
+from typing import Optional
+from loguru import logger
 
 import pytest
 from fastapi import WebSocket
@@ -16,6 +17,8 @@ fixtures = get_fixtures("clients")
 alice = fixtures["alice"]
 bob = fixtures["bob"]
 
+RELAY_ID = "relay_01"
+
 
 class MockWebSocket(WebSocket):
     def __init__(self):
@@ -36,10 +39,15 @@ class MockWebSocket(WebSocket):
     async def wire_mock_data(self, data: dict):
         await self.fake_wire.put(dumps(data))
 
+    async def close(
+        self, code: int = 1000, reason: Optional[str] = None
+    ) -> None:
+        logger.info(reason)
+
 
 @pytest.mark.asyncio
 async def test_alice_and_bob():
-    ws_alice, ws_bob = init_clients()
+    ws_alice, ws_bob = await init_clients()
 
     await alice_wires_meta_and_post01(ws_alice)
 
@@ -62,17 +70,19 @@ async def test_alice_and_bob():
     await alice_deletes_post01__bob_is_notified(ws_alice, ws_bob)
 
 
-def init_clients():
+async def init_clients():
     client_manager = NostrClientManager()
+    client_manager.active_relays = [RELAY_ID]
+    client_manager.toggle_relay(RELAY_ID, True)
 
     ws_alice = MockWebSocket()
-    client_alice = NostrClientConnection(websocket=ws_alice)
-    client_manager.add_client(client_alice)
+    client_alice = NostrClientConnection(relay_id=RELAY_ID, websocket=ws_alice)
+    await client_manager.add_client(client_alice)
     asyncio.create_task(client_alice.start())
 
     ws_bob = MockWebSocket()
-    client_bob = NostrClientConnection(websocket=ws_bob)
-    client_manager.add_client(client_bob)
+    client_bob = NostrClientConnection(relay_id=RELAY_ID, websocket=ws_bob)
+    await client_manager.add_client(client_bob)
     asyncio.create_task(client_bob.start())
     return ws_alice, ws_bob