package apis_test import ( "errors" "net/http" "strings" "testing" "time" "github.com/labstack/echo/v5" "github.com/pocketbase/dbx" "github.com/pocketbase/pocketbase/apis" "github.com/pocketbase/pocketbase/core" "github.com/pocketbase/pocketbase/daos" "github.com/pocketbase/pocketbase/models" "github.com/pocketbase/pocketbase/tests" "github.com/pocketbase/pocketbase/tools/hook" "github.com/pocketbase/pocketbase/tools/subscriptions" ) func TestRealtimeConnect(t *testing.T) { scenarios := []tests.ApiScenario{ { Method: http.MethodGet, Url: "/api/realtime", Timeout: 100 * time.Millisecond, ExpectedStatus: 200, ExpectedContent: []string{ `id:`, `event:PB_CONNECT`, `data:{"clientId":`, }, ExpectedEvents: map[string]int{ "OnRealtimeConnectRequest": 1, "OnRealtimeBeforeMessageSend": 1, "OnRealtimeAfterMessageSend": 1, "OnRealtimeDisconnectRequest": 1, }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { if len(app.SubscriptionsBroker().Clients()) != 0 { t.Errorf("Expected the subscribers to be removed after connection close, found %d", len(app.SubscriptionsBroker().Clients())) } }, }, { Name: "PB_CONNECT interrupt", Method: http.MethodGet, Url: "/api/realtime", Timeout: 100 * time.Millisecond, ExpectedStatus: 200, ExpectedEvents: map[string]int{ "OnRealtimeConnectRequest": 1, "OnRealtimeBeforeMessageSend": 1, "OnRealtimeDisconnectRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { app.OnRealtimeBeforeMessageSend().Add(func(e *core.RealtimeMessageEvent) error { if e.Message.Name == "PB_CONNECT" { return errors.New("PB_CONNECT error") } return nil }) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { if len(app.SubscriptionsBroker().Clients()) != 0 { t.Errorf("Expected the subscribers to be removed after connection close, found %d", len(app.SubscriptionsBroker().Clients())) } }, }, { Name: "Skipping/ignoring messages", Method: http.MethodGet, Url: "/api/realtime", Timeout: 100 * time.Millisecond, ExpectedStatus: 200, ExpectedEvents: map[string]int{ "OnRealtimeConnectRequest": 1, "OnRealtimeBeforeMessageSend": 1, "OnRealtimeDisconnectRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { app.OnRealtimeBeforeMessageSend().Add(func(e *core.RealtimeMessageEvent) error { return hook.StopPropagation }) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { if len(app.SubscriptionsBroker().Clients()) != 0 { t.Errorf("Expected the subscribers to be removed after connection close, found %d", len(app.SubscriptionsBroker().Clients())) } }, }, } for _, scenario := range scenarios { scenario.Test(t) } } func TestRealtimeSubscribe(t *testing.T) { client := subscriptions.NewDefaultClient() resetClient := func() { client.Unsubscribe() client.Set(apis.ContextAdminKey, nil) client.Set(apis.ContextAuthRecordKey, nil) } scenarios := []tests.ApiScenario{ { Name: "missing client", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"missing","subscriptions":["test1", "test2"]}`), ExpectedStatus: 404, ExpectedContent: []string{`"data":{}`}, }, { Name: "existing client - empty subscriptions", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"` + client.Id() + `","subscriptions":[]}`), ExpectedStatus: 204, ExpectedEvents: map[string]int{ "OnRealtimeBeforeSubscribeRequest": 1, "OnRealtimeAfterSubscribeRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { client.Subscribe("test0") app.SubscriptionsBroker().Register(client) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { if len(client.Subscriptions()) != 0 { t.Errorf("Expected no subscriptions, got %v", client.Subscriptions()) } resetClient() }, }, { Name: "existing client - 2 new subscriptions", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"` + client.Id() + `","subscriptions":["test1", "test2"]}`), ExpectedStatus: 204, ExpectedEvents: map[string]int{ "OnRealtimeBeforeSubscribeRequest": 1, "OnRealtimeAfterSubscribeRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { client.Subscribe("test0") app.SubscriptionsBroker().Register(client) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { expectedSubs := []string{"test1", "test2"} if len(expectedSubs) != len(client.Subscriptions()) { t.Errorf("Expected subscriptions %v, got %v", expectedSubs, client.Subscriptions()) } for _, s := range expectedSubs { if !client.HasSubscription(s) { t.Errorf("Cannot find %q subscription in %v", s, client.Subscriptions()) } } resetClient() }, }, { Name: "existing client - authorized admin", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"` + client.Id() + `","subscriptions":["test1", "test2"]}`), RequestHeaders: map[string]string{ "Authorization": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpZCI6InN5d2JoZWNuaDQ2cmhtMCIsInR5cGUiOiJhZG1pbiIsImV4cCI6MjIwODk4NTI2MX0.M1m--VOqGyv0d23eeUc0r9xE8ZzHaYVmVFw1VZW6gT8", }, ExpectedStatus: 204, ExpectedEvents: map[string]int{ "OnRealtimeBeforeSubscribeRequest": 1, "OnRealtimeAfterSubscribeRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { app.SubscriptionsBroker().Register(client) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { admin, _ := client.Get(apis.ContextAdminKey).(*models.Admin) if admin == nil { t.Errorf("Expected admin auth model, got nil") } resetClient() }, }, { Name: "existing client - authorized record", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"` + client.Id() + `","subscriptions":["test1", "test2"]}`), RequestHeaders: map[string]string{ "Authorization": "eyJhbGciOiJIUzI1NiJ9.eyJpZCI6IjRxMXhsY2xtZmxva3UzMyIsInR5cGUiOiJhdXRoUmVjb3JkIiwiY29sbGVjdGlvbklkIjoiX3BiX3VzZXJzX2F1dGhfIiwiZXhwIjoyMjA4OTg1MjYxfQ.UwD8JvkbQtXpymT09d7J6fdA0aP9g4FJ1GPh_ggEkzc", }, ExpectedStatus: 204, ExpectedEvents: map[string]int{ "OnRealtimeBeforeSubscribeRequest": 1, "OnRealtimeAfterSubscribeRequest": 1, }, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { app.SubscriptionsBroker().Register(client) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { authRecord, _ := client.Get(apis.ContextAuthRecordKey).(*models.Record) if authRecord == nil { t.Errorf("Expected auth record model, got nil") } resetClient() }, }, { Name: "existing client - mismatched auth", Method: http.MethodPost, Url: "/api/realtime", Body: strings.NewReader(`{"clientId":"` + client.Id() + `","subscriptions":["test1", "test2"]}`), RequestHeaders: map[string]string{ "Authorization": "eyJhbGciOiJIUzI1NiJ9.eyJpZCI6IjRxMXhsY2xtZmxva3UzMyIsInR5cGUiOiJhdXRoUmVjb3JkIiwiY29sbGVjdGlvbklkIjoiX3BiX3VzZXJzX2F1dGhfIiwiZXhwIjoyMjA4OTg1MjYxfQ.UwD8JvkbQtXpymT09d7J6fdA0aP9g4FJ1GPh_ggEkzc", }, ExpectedStatus: 403, ExpectedContent: []string{`"data":{}`}, BeforeTestFunc: func(t *testing.T, app *tests.TestApp, e *echo.Echo) { initialAuth := &models.Record{} initialAuth.RefreshId() client.Set(apis.ContextAuthRecordKey, initialAuth) app.SubscriptionsBroker().Register(client) }, AfterTestFunc: func(t *testing.T, app *tests.TestApp, res *http.Response) { authRecord, _ := client.Get(apis.ContextAuthRecordKey).(*models.Record) if authRecord == nil { t.Errorf("Expected auth record model, got nil") } resetClient() }, }, } for _, scenario := range scenarios { scenario.Test(t) } } func TestRealtimeAuthRecordDeleteEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) authRecord1, err := testApp.Dao().FindFirstRecordByData("users", "email", "test@example.com") if err != nil { t.Fatal(err) } authRecord2, err := testApp.Dao().FindFirstRecordByData("users", "email", "test2@example.com") if err != nil { t.Fatal(err) } client1 := subscriptions.NewDefaultClient() client1.Set(apis.ContextAuthRecordKey, authRecord1) testApp.SubscriptionsBroker().Register(client1) client2 := subscriptions.NewDefaultClient() client2.Set(apis.ContextAuthRecordKey, authRecord1) testApp.SubscriptionsBroker().Register(client2) client3 := subscriptions.NewDefaultClient() client3.Set(apis.ContextAuthRecordKey, authRecord2) testApp.SubscriptionsBroker().Register(client3) // "delete" authRecord1 e := new(core.ModelEvent) e.Dao = testApp.Dao() e.Model = authRecord1 testApp.OnModelAfterDelete().Trigger(e) if total := len(testApp.SubscriptionsBroker().Clients()); total != 3 { t.Fatalf("Expected %d subscription clients, found %d", 3, total) } if auth := client1.Get(apis.ContextAuthRecordKey); auth != nil { t.Fatalf("[client1] Expected the auth state to be unset, found %#v", auth) } if auth := client2.Get(apis.ContextAuthRecordKey); auth != nil { t.Fatalf("[client2] Expected the auth state to be unset, found %#v", auth) } if auth := client3.Get(apis.ContextAuthRecordKey); auth == nil || auth.(*models.Record).Id != authRecord2.Id { t.Fatalf("[client3] Expected the auth state to be left unchanged, found %#v", auth) } } func TestRealtimeAuthRecordUpdateEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) authRecord1, err := testApp.Dao().FindFirstRecordByData("users", "email", "test@example.com") if err != nil { t.Fatal(err) } client := subscriptions.NewDefaultClient() client.Set(apis.ContextAuthRecordKey, authRecord1) testApp.SubscriptionsBroker().Register(client) // refetch the authRecord and change its email authRecord2, err := testApp.Dao().FindFirstRecordByData("users", "email", "test@example.com") if err != nil { t.Fatal(err) } authRecord2.SetEmail("new@example.com") e := new(core.ModelEvent) e.Dao = testApp.Dao() e.Model = authRecord2 testApp.OnModelAfterUpdate().Trigger(e) clientAuthRecord, _ := client.Get(apis.ContextAuthRecordKey).(*models.Record) if clientAuthRecord.Email() != authRecord2.Email() { t.Fatalf("Expected authRecord with email %q, got %q", authRecord2.Email(), clientAuthRecord.Email()) } } func TestRealtimeAdminDeleteEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) admin1, err := testApp.Dao().FindAdminByEmail("test@example.com") if err != nil { t.Fatal(err) } admin2, err := testApp.Dao().FindAdminByEmail("test2@example.com") if err != nil { t.Fatal(err) } client1 := subscriptions.NewDefaultClient() client1.Set(apis.ContextAdminKey, admin1) testApp.SubscriptionsBroker().Register(client1) client2 := subscriptions.NewDefaultClient() client2.Set(apis.ContextAdminKey, admin1) testApp.SubscriptionsBroker().Register(client2) client3 := subscriptions.NewDefaultClient() client3.Set(apis.ContextAdminKey, admin2) testApp.SubscriptionsBroker().Register(client3) // "delete" admin1 e := new(core.ModelEvent) e.Dao = testApp.Dao() e.Model = admin1 testApp.OnModelAfterDelete().Trigger(e) if total := len(testApp.SubscriptionsBroker().Clients()); total != 3 { t.Fatalf("Expected %d subscription clients, found %d", 3, total) } if auth := client1.Get(apis.ContextAdminKey); auth != nil { t.Fatalf("[client1] Expected the auth state to be unset, found %#v", auth) } if auth := client2.Get(apis.ContextAdminKey); auth != nil { t.Fatalf("[client2] Expected the auth state to be unset, found %#v", auth) } if auth := client3.Get(apis.ContextAdminKey); auth == nil || auth.(*models.Admin).Id != admin2.Id { t.Fatalf("[client3] Expected the auth state to be left unchanged, found %#v", auth) } } func TestRealtimeAdminUpdateEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) admin1, err := testApp.Dao().FindAdminByEmail("test@example.com") if err != nil { t.Fatal(err) } client := subscriptions.NewDefaultClient() client.Set(apis.ContextAdminKey, admin1) testApp.SubscriptionsBroker().Register(client) // refetch the authRecord and change its email admin2, err := testApp.Dao().FindAdminByEmail("test@example.com") if err != nil { t.Fatal(err) } admin2.Email = "new@example.com" e := new(core.ModelEvent) e.Dao = testApp.Dao() e.Model = admin2 testApp.OnModelAfterUpdate().Trigger(e) clientAdmin, _ := client.Get(apis.ContextAdminKey).(*models.Admin) if clientAdmin.Email != admin2.Email { t.Fatalf("Expected authRecord with email %q, got %q", admin2.Email, clientAdmin.Email) } } // Custom auth record model struct // ------------------------------------------------------------------- var _ models.Model = (*CustomUser)(nil) type CustomUser struct { models.BaseModel Email string `db:"email" json:"email"` } func (m *CustomUser) TableName() string { return "users" // the name of your collection } func findCustomUserByEmail(dao *daos.Dao, email string) (*CustomUser, error) { model := &CustomUser{} err := dao.ModelQuery(model). AndWhere(dbx.HashExp{"email": email}). Limit(1). One(model) if err != nil { return nil, err } return model, nil } func TestRealtimeCustomAuthModelDeleteEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) authRecord1, err := testApp.Dao().FindFirstRecordByData("users", "email", "test@example.com") if err != nil { t.Fatal(err) } authRecord2, err := testApp.Dao().FindFirstRecordByData("users", "email", "test2@example.com") if err != nil { t.Fatal(err) } client1 := subscriptions.NewDefaultClient() client1.Set(apis.ContextAuthRecordKey, authRecord1) testApp.SubscriptionsBroker().Register(client1) client2 := subscriptions.NewDefaultClient() client2.Set(apis.ContextAuthRecordKey, authRecord1) testApp.SubscriptionsBroker().Register(client2) client3 := subscriptions.NewDefaultClient() client3.Set(apis.ContextAuthRecordKey, authRecord2) testApp.SubscriptionsBroker().Register(client3) // refetch authRecord1 as CustomUser customUser, err := findCustomUserByEmail(testApp.Dao(), "test@example.com") if err != nil { t.Fatal(err) } // delete the custom user (should unset the client auth record) if err := testApp.Dao().Delete(customUser); err != nil { t.Fatal(err) } if total := len(testApp.SubscriptionsBroker().Clients()); total != 3 { t.Fatalf("Expected %d subscription clients, found %d", 3, total) } if auth := client1.Get(apis.ContextAuthRecordKey); auth != nil { t.Fatalf("[client1] Expected the auth state to be unset, found %#v", auth) } if auth := client2.Get(apis.ContextAuthRecordKey); auth != nil { t.Fatalf("[client2] Expected the auth state to be unset, found %#v", auth) } if auth := client3.Get(apis.ContextAuthRecordKey); auth == nil || auth.(*models.Record).Id != authRecord2.Id { t.Fatalf("[client3] Expected the auth state to be left unchanged, found %#v", auth) } } func TestRealtimeCustomAuthModelUpdateEvent(t *testing.T) { testApp, _ := tests.NewTestApp() defer testApp.Cleanup() apis.InitApi(testApp) authRecord, err := testApp.Dao().FindFirstRecordByData("users", "email", "test@example.com") if err != nil { t.Fatal(err) } client := subscriptions.NewDefaultClient() client.Set(apis.ContextAuthRecordKey, authRecord) testApp.SubscriptionsBroker().Register(client) // refetch the authRecord as CustomUser customUser, err := findCustomUserByEmail(testApp.Dao(), "test@example.com") if err != nil { t.Fatal(err) } // change its email customUser.Email = "new@example.com" if err := testApp.Dao().Save(customUser); err != nil { t.Fatal(err) } clientAuthRecord, _ := client.Get(apis.ContextAuthRecordKey).(*models.Record) if clientAuthRecord.Email() != customUser.Email { t.Fatalf("Expected authRecord with email %q, got %q", customUser.Email, clientAuthRecord.Email()) } }