fix: resolve flakey inbox tests (coder#17010)

Author: DanielleMaywood

coderd/inboxnotifications.go

@@ -94,18 +94,6 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	conn, err := websocket.Accept(rw, r, nil)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Failed to upgrade connection to websocket.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-
-	go httpapi.Heartbeat(ctx, conn)
-	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
-
 	notificationCh := make(chan codersdk.InboxNotification, 10)
 
 	closeInboxNotificationsSubscriber, err := api.Pubsub.SubscribeWithErr(pubsub.InboxNotificationForOwnerEventChannel(apikey.UserID),
@@ -161,9 +149,20 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request)
 		api.Logger.Error(ctx, "subscribe to inbox notification event", slog.Error(err))
 		return
 	}
-
 	defer closeInboxNotificationsSubscriber()
 
+	conn, err := websocket.Accept(rw, r, nil)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to upgrade connection to websocket.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	go httpapi.Heartbeat(ctx, conn)
+	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
+
 	encoder := wsjson.NewEncoder[codersdk.GetInboxNotificationResponse](conn, websocket.MessageText)
 	defer encoder.Close(websocket.StatusNormalClosure)
 

coderd/inboxnotifications_test.go

@@ -122,7 +122,8 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "notification title", "notification content", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		_, message, err := wsConn.Read(ctx)
 		require.NoError(t, err)
@@ -174,7 +175,8 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "memory related title", "memory related content", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		_, message, err := wsConn.Read(ctx)
 		require.NoError(t, err)
@@ -193,15 +195,17 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "disk related title", "disk related title", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		dispatchFunc, err = inboxHandler.Dispatcher(types.MessagePayload{
 			UserID:                 memberClient.ID.String(),
 			NotificationTemplateID: notifications.TemplateWorkspaceOutOfMemory.String(),
 		}, "second memory related title", "second memory related title", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		_, message, err = wsConn.Read(ctx)
 		require.NoError(t, err)
@@ -256,7 +260,8 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "memory related title", "memory related content", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		_, message, err := wsConn.Read(ctx)
 		require.NoError(t, err)
@@ -276,7 +281,8 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "second memory related title", "second memory related title", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		dispatchFunc, err = inboxHandler.Dispatcher(types.MessagePayload{
 			UserID:                 memberClient.ID.String(),
@@ -285,7 +291,8 @@ func TestInboxNotification_Watch(t *testing.T) {
 		}, "another memory related title", "another memory related title", nil)
 		require.NoError(t, err)
 
-		dispatchFunc(ctx, uuid.New())
+		_, err = dispatchFunc(ctx, uuid.New())
+		require.NoError(t, err)
 
 		_, message, err = wsConn.Read(ctx)
 		require.NoError(t, err)