diff --git a/testpds/testpds.go b/testpds/testpds.go index cd524bd..348b742 100644 --- a/testpds/testpds.go +++ b/testpds/testpds.go @@ -10,6 +10,7 @@ import ( "crypto/rand" "encoding/json" "fmt" + "log/slog" "net" "net/http" "testing" @@ -21,7 +22,10 @@ import ( "github.com/bluesky-social/indigo/api/atproto" "github.com/bluesky-social/indigo/atproto/atcrypto" + "github.com/bluesky-social/indigo/events" + "github.com/bluesky-social/indigo/events/schedulers/sequential" "github.com/bluesky-social/indigo/xrpc" + "github.com/gorilla/websocket" cocoonplc "github.com/haileyok/cocoon/plc" "github.com/haileyok/cocoon/server" "github.com/lestrrat-go/jwx/v2/jwk" @@ -201,6 +205,67 @@ func (p *TestPDS) MustCreateAccount(t *testing.T, handle, email, password string }) } +// FirehoseEvent is a union of the event types delivered by subscribeRepos. +// Exactly one field will be non-nil per event. +type FirehoseEvent struct { + Commit *atproto.SyncSubscribeRepos_Commit + Identity *atproto.SyncSubscribeRepos_Identity + Account *atproto.SyncSubscribeRepos_Account +} + +// Subscribe connects to the PDS firehose (com.atproto.sync.subscribeRepos) and +// delivers events on the returned channel. The connection is closed and the +// channel drained when ctx is cancelled. An optional cursor resumes from a +// specific sequence number (pass 0 to start from live). +func (p *TestPDS) Subscribe(ctx context.Context, cursor int64) (<-chan FirehoseEvent, error) { + wsURL := fmt.Sprintf("ws://%s/xrpc/com.atproto.sync.subscribeRepos", p.URL[len("http://"):]) + if cursor > 0 { + wsURL = fmt.Sprintf("%s?cursor=%d", wsURL, cursor) + } + + conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, nil) + if err != nil { + return nil, fmt.Errorf("testpds: subscribe: %w", err) + } + + ch := make(chan FirehoseEvent, 256) + + rsc := &events.RepoStreamCallbacks{ + RepoCommit: func(evt *atproto.SyncSubscribeRepos_Commit) error { + select { + case ch <- FirehoseEvent{Commit: evt}: + case <-ctx.Done(): + } + return nil + }, + RepoIdentity: func(evt *atproto.SyncSubscribeRepos_Identity) error { + select { + case ch <- FirehoseEvent{Identity: evt}: + case <-ctx.Done(): + } + return nil + }, + RepoAccount: func(evt *atproto.SyncSubscribeRepos_Account) error { + select { + case ch <- FirehoseEvent{Account: evt}: + case <-ctx.Done(): + } + return nil + }, + } + + sched := sequential.NewScheduler("testpds", rsc.EventHandler) + + go func() { + defer close(ch) + defer conn.Close() + defer sched.Shutdown() + _ = events.HandleRepoStream(ctx, conn, sched, slog.Default()) + }() + + return ch, nil +} + func freePort(t *testing.T) int { t.Helper() l, err := net.Listen("tcp", "127.0.0.1:0") diff --git a/testpds/testpds_test.go b/testpds/testpds_test.go index c7d6ee8..01300d3 100644 --- a/testpds/testpds_test.go +++ b/testpds/testpds_test.go @@ -96,6 +96,48 @@ func TestMustCreateAccount(t *testing.T) { } } +func TestSubscribe(t *testing.T) { + pds := testpds.Start(t, nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ch, err := pds.Subscribe(ctx, 0) + if err != nil { + t.Fatalf("subscribe: %v", err) + } + + // Create an account — this should emit identity + commit events + alice := pds.MustCreateAccount(t, "alice.test", "alice@test.com", "password123") + + // Create a record to generate a commit event + _, err = atproto.RepoCreateRecord(ctx, alice, &atproto.RepoCreateRecord_Input{ + Repo: alice.Auth.Did, + Collection: "app.bsky.feed.post", + Record: &lexutil.LexiconTypeDecoder{Val: &bsky.FeedPost{ + Text: "hello firehose", + CreatedAt: time.Now().Format(time.RFC3339), + }}, + }) + if err != nil { + t.Fatalf("create record: %v", err) + } + + // We should receive at least one commit event for the post + deadline := time.After(5 * time.Second) + var gotCommit bool + for !gotCommit { + select { + case evt := <-ch: + if evt.Commit != nil && evt.Commit.Repo == alice.Auth.Did { + gotCommit = true + t.Logf("got commit event: repo=%s seq=%d ops=%d", evt.Commit.Repo, evt.Commit.Seq, len(evt.Commit.Ops)) + } + case <-deadline: + t.Fatal("timed out waiting for commit event") + } + } +} + func TestRecordCRUD(t *testing.T) { pds := testpds.Start(t, nil) ctx := context.Background()