1 files changed, 333 insertions, 0 deletions

diff --git a/relay_test.go b/relay_test.go
new file mode 100644
index 0000000..4ace956
--- /dev/null
+++ b/relay_test.go
@@ -0,0 +1,333 @@
+package nostr
+
+import (
+        "context"
+        "encoding/json"
+        "net/http"
+        "net/http/httptest"
+        "strings"
+        "testing"
+        "time"
+
+        "github.com/coder/websocket"
+)
+
+// mockRelay creates a test WebSocket server that echoes messages
+func mockRelay(t *testing.T, handler func(conn *websocket.Conn)) *httptest.Server {
+        return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+                conn, err := websocket.Accept(w, r, nil)
+                if err != nil {
+                        t.Logf("Failed to accept WebSocket: %v", err)
+                        return
+                }
+                defer conn.Close(websocket.StatusNormalClosure, "")
+
+                handler(conn)
+        }))
+}
+
+func TestConnect(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Just accept and wait
+                time.Sleep(100 * time.Millisecond)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        if relay.URL != url {
+                t.Errorf("Relay.URL = %s, want %s", relay.URL, url)
+        }
+}
+
+func TestConnectError(t *testing.T) {
+        ctx := context.Background()
+        _, err := Connect(ctx, "ws://localhost:99999")
+        if err == nil {
+                t.Error("Connect() expected error for invalid URL")
+        }
+}
+
+func TestRelaySendReceive(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Read message
+                _, data, err := conn.Read(context.Background())
+                if err != nil {
+                        t.Logf("Read error: %v", err)
+                        return
+                }
+
+                // Echo it back as NOTICE
+                var arr []interface{}
+                json.Unmarshal(data, &arr)
+
+                response, _ := json.Marshal([]interface{}{"NOTICE", "received: " + arr[0].(string)})
+                conn.Write(context.Background(), websocket.MessageText, response)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        // Send a CLOSE envelope
+        closeEnv := CloseEnvelope{SubscriptionID: "test"}
+        if err := relay.Send(ctx, closeEnv); err != nil {
+                t.Fatalf("Send() error = %v", err)
+        }
+
+        // Receive response
+        env, err := relay.Receive(ctx)
+        if err != nil {
+                t.Fatalf("Receive() error = %v", err)
+        }
+
+        noticeEnv, ok := env.(*NoticeEnvelope)
+        if !ok {
+                t.Fatalf("Expected *NoticeEnvelope, got %T", env)
+        }
+
+        if !strings.Contains(noticeEnv.Message, "CLOSE") {
+                t.Errorf("Message = %s, want to contain 'CLOSE'", noticeEnv.Message)
+        }
+}
+
+func TestRelayPublish(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Read the EVENT message
+                _, data, err := conn.Read(context.Background())
+                if err != nil {
+                        t.Logf("Read error: %v", err)
+                        return
+                }
+
+                // Parse to get event ID
+                var arr []json.RawMessage
+                json.Unmarshal(data, &arr)
+
+                var event Event
+                json.Unmarshal(arr[1], &event)
+
+                // Send OK response
+                response, _ := json.Marshal([]interface{}{"OK", event.ID, true, ""})
+                conn.Write(context.Background(), websocket.MessageText, response)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        // Create and sign event
+        key, _ := GenerateKey()
+        event := &Event{
+                CreatedAt: time.Now().Unix(),
+                Kind:      KindTextNote,
+                Tags:      Tags{},
+                Content:   "Test publish",
+        }
+        key.Sign(event)
+
+        // Publish
+        if err := relay.Publish(ctx, event); err != nil {
+                t.Fatalf("Publish() error = %v", err)
+        }
+}
+
+func TestRelayPublishRejected(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Read the EVENT message
+                _, data, err := conn.Read(context.Background())
+                if err != nil {
+                        return
+                }
+
+                var arr []json.RawMessage
+                json.Unmarshal(data, &arr)
+
+                var event Event
+                json.Unmarshal(arr[1], &event)
+
+                // Send rejection
+                response, _ := json.Marshal([]interface{}{"OK", event.ID, false, "blocked: spam"})
+                conn.Write(context.Background(), websocket.MessageText, response)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        key, _ := GenerateKey()
+        event := &Event{
+                CreatedAt: time.Now().Unix(),
+                Kind:      KindTextNote,
+                Tags:      Tags{},
+                Content:   "Test",
+        }
+        key.Sign(event)
+
+        err = relay.Publish(ctx, event)
+        if err == nil {
+                t.Error("Publish() expected error for rejected event")
+        }
+        if !strings.Contains(err.Error(), "rejected") {
+                t.Errorf("Error = %v, want to contain 'rejected'", err)
+        }
+}
+
+func TestRelaySubscribe(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Read REQ
+                _, data, err := conn.Read(context.Background())
+                if err != nil {
+                        return
+                }
+
+                var arr []json.RawMessage
+                json.Unmarshal(data, &arr)
+
+                var subID string
+                json.Unmarshal(arr[1], &subID)
+
+                // Send some events
+                for i := 0; i < 3; i++ {
+                        event := Event{
+                                ID:        "event" + string(rune('0'+i)),
+                                PubKey:    "pubkey",
+                                CreatedAt: time.Now().Unix(),
+                                Kind:      1,
+                                Tags:      Tags{},
+                                Content:   "Test event",
+                                Sig:       "sig",
+                        }
+                        response, _ := json.Marshal([]interface{}{"EVENT", subID, event})
+                        conn.Write(context.Background(), websocket.MessageText, response)
+                }
+
+                // Send EOSE
+                eose, _ := json.Marshal([]interface{}{"EOSE", subID})
+                conn.Write(context.Background(), websocket.MessageText, eose)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        sub, err := relay.Subscribe(ctx, "sub1", Filter{Kinds: []int{1}})
+        if err != nil {
+                t.Fatalf("Subscribe() error = %v", err)
+        }
+
+        // Start listening in background
+        go relay.Listen(ctx)
+
+        // Collect events
+        eventCount := 0
+        timeout := time.After(2 * time.Second)
+
+        for {
+                select {
+                case <-sub.Events:
+                        eventCount++
+                case <-sub.EOSE:
+                        if eventCount != 3 {
+                                t.Errorf("Received %d events, want 3", eventCount)
+                        }
+                        return
+                case <-timeout:
+                        t.Fatal("Timeout waiting for events")
+                }
+        }
+}
+
+func TestRelayClose(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                time.Sleep(100 * time.Millisecond)
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+
+        if err := relay.Close(); err != nil {
+                t.Errorf("Close() error = %v", err)
+        }
+
+        // Second close should be safe
+        if err := relay.Close(); err != nil {
+                t.Errorf("Second Close() error = %v", err)
+        }
+}
+
+func TestSubscriptionClose(t *testing.T) {
+        server := mockRelay(t, func(conn *websocket.Conn) {
+                // Read REQ
+                conn.Read(context.Background())
+
+                // Wait for CLOSE
+                _, data, err := conn.Read(context.Background())
+                if err != nil {
+                        return
+                }
+
+                var arr []interface{}
+                json.Unmarshal(data, &arr)
+
+                if arr[0] != "CLOSE" {
+                        t.Errorf("Expected CLOSE, got %v", arr[0])
+                }
+        })
+        defer server.Close()
+
+        url := "ws" + strings.TrimPrefix(server.URL, "http")
+        ctx := context.Background()
+
+        relay, err := Connect(ctx, url)
+        if err != nil {
+                t.Fatalf("Connect() error = %v", err)
+        }
+        defer relay.Close()
+
+        sub, err := relay.Subscribe(ctx, "sub1", Filter{Kinds: []int{1}})
+        if err != nil {
+                t.Fatalf("Subscribe() error = %v", err)
+        }
+
+        if err := sub.Close(ctx); err != nil {
+                t.Errorf("Subscription.Close() error = %v", err)
+        }
+}