Function TestSubscribeRelaySnapshotDelivered

enterprise/coderd/x/chatd/chatd_test.go:401–485

(t *testing.T)

}

func TestSubscribeRelaySnapshotDelivered(t *testing.T) {
	t.Parallel()

	db, ps := dbtestutil.NewDB(t)
	workerID := uuid.New()
	subscriberID := uuid.New()

	provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
		[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
	) {
		// Return a non-empty snapshot with two parts.
		snapshot := []codersdk.ChatStreamEvent{
			{
				Type: codersdk.ChatStreamEventTypeMessagePart,
				MessagePart: &codersdk.ChatStreamMessagePart{
					Role: "assistant",
					Part: codersdk.ChatMessageText("snap-one"),
				},
			},
			{
				Type: codersdk.ChatStreamEventTypeMessagePart,
				MessagePart: &codersdk.ChatStreamMessagePart{
					Role: "assistant",
					Part: codersdk.ChatMessageText("snap-two"),
				},
			},
		}
		ch := make(chan codersdk.ChatStreamEvent, 10)
		// Also send a live part after the snapshot.
		ch <- codersdk.ChatStreamEvent{
			Type: codersdk.ChatStreamEventTypeMessagePart,
			MessagePart: &codersdk.ChatStreamMessagePart{
				Role: "assistant",
				Part: codersdk.ChatMessageText("live-part"),
			},
		}
		return snapshot, ch, func() {}, nil
	}

	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)

	ctx := testutil.Context(t, testutil.WaitLong)
	user, org, model := seedChatDependencies(t, db)

	chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-snapshot")
	staleChat := chat
	staleChat.Status = database.ChatStatusWaiting
	staleChat.WorkerID = uuid.NullUUID{}
	staleChat.StartedAt = sql.NullTime{}
	staleChat.HeartbeatAt = sql.NullTime{}

	initialSnapshot, events, cancel, ok := subscriber.SubscribeAuthorized(ctx, staleChat, nil, 0)
	require.True(t, ok)
	t.Cleanup(cancel)

	// The relay snapshot parts are forwarded through the events
	// channel by the enterprise SubscribeFn. Collect them along
	// with the live part.