From af049a258a90e924a2439cdaae5dd08ce4447a02 Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Sun, 6 Jul 2025 17:07:34 -0400
Subject: [PATCH] fix relay health check

---
 src/nostr/client.py            |  4 +---
 src/tests/test_nostr_client.py | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/nostr/client.py b/src/nostr/client.py
index 2ff7019..320ba72 100644
--- a/src/nostr/client.py
+++ b/src/nostr/client.py
@@ -153,10 +153,8 @@ class NostrClient:
                 while True:
                     msg = await asyncio.wait_for(ws.recv(), timeout=timeout)
                     data = json.loads(msg)
-                    if data[0] == "EVENT":
+                    if data[0] in {"EVENT", "EOSE"}:
                         return True
-                    if data[0] == "EOSE":
-                        return False
         except Exception:
             return False
 
diff --git a/src/tests/test_nostr_client.py b/src/tests/test_nostr_client.py
index 8a849b8..310ab4b 100644
--- a/src/tests/test_nostr_client.py
+++ b/src/tests/test_nostr_client.py
@@ -2,12 +2,15 @@ import sys
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest.mock import patch
+import json
+import asyncio
 from cryptography.fernet import Fernet
 
 sys.path.append(str(Path(__file__).resolve().parents[1]))
 
 from password_manager.encryption import EncryptionManager
 from nostr.client import NostrClient
+import nostr.client as nostr_client
 
 
 def test_nostr_client_uses_custom_relays():
@@ -50,6 +53,25 @@ class FakeAddRelayClient:
         self.connected = True
 
 
+class FakeWebSocket:
+    def __init__(self, messages):
+        self.messages = messages
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+    async def send(self, _):
+        pass
+
+    async def recv(self):
+        if self.messages:
+            return self.messages.pop(0)
+        await asyncio.sleep(0)
+
+
 def _setup_client(tmpdir, fake_cls):
     key = Fernet.generate_key()
     enc_mgr = EncryptionManager(key, Path(tmpdir))
@@ -91,3 +113,18 @@ def test_check_relay_health_runs_async(tmp_path, monkeypatch):
 
     assert result == 1
     assert recorded["args"] == (3, 2)
+
+
+def test_ping_relay_accepts_eose(tmp_path, monkeypatch):
+    client = _setup_client(tmp_path, FakeAddRelayClient)
+
+    fake_ws = FakeWebSocket([json.dumps(["EOSE"])])
+
+    def fake_connect(*_args, **_kwargs):
+        return fake_ws
+
+    monkeypatch.setattr(nostr_client.websockets, "connect", fake_connect)
+
+    result = asyncio.run(client._ping_relay("wss://relay", timeout=0.1))
+
+    assert result is True