path: root/relay_test.go

package nostr

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"code.northwest.io/nostr/internal/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)
		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()

	// Create relay without auto-Listen to test Send/Receive directly
	conn, err := websocket.Dial(ctx, url)
	if err != nil {
		t.Fatalf("Dial() error = %v", err)
	}
	relay := &Relay{
		URL:           url,
		conn:          conn,
		subscriptions: make(map[string]*Subscription),
		okChannels:    make(map[string]chan *OKEnvelope),
	}
	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, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	relay, err := Connect(ctx, url)
	if err != nil {
		t.Fatalf("Connect() error = %v", err)
	}
	defer relay.Close()

	sub := relay.Fetch(ctx, Filter{Kinds: []int{1}})

	eventCount := 0
	for range sub.Events {
		eventCount++
	}

	if eventCount != 3 {
		t.Errorf("Received %d events, want 3", eventCount)
	}
	if sub.Err != nil {
		t.Errorf("Subscription.Err = %v, want nil", sub.Err)
	}
}

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 := relay.Subscribe(ctx, Filter{Kinds: []int{1}})

	if err := sub.Close(ctx); err != nil {
		t.Errorf("Subscription.Close() error = %v", err)
	}
}