Function TestUpdateChatHeartbeatsRequiresOwnership

coderd/x/chatd/chatd_test.go:1835–1881

(t *testing.T)

}

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

	db, ps := dbtestutil.NewDB(t)
	replica := newTestServer(t, db, ps, uuid.New())

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

	chat, err := replica.CreateChat(ctx, chatd.CreateOptions{
		OrganizationID:     org.ID,
		OwnerID:            user.ID,
		Title:              "heartbeat-ownership",
		ModelConfigID:      model.ID,
		InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
	})
	require.NoError(t, err)

	workerID := uuid.New()
	chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
		ID:          chat.ID,
		Status:      database.ChatStatusRunning,
		WorkerID:    uuid.NullUUID{UUID: workerID, Valid: true},
		StartedAt:   sql.NullTime{Time: time.Now(), Valid: true},
		HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true},
	})
	require.NoError(t, err)

	// Wrong worker_id should return no IDs.
	ids, err := db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{
		IDs:      []uuid.UUID{chat.ID},
		WorkerID: uuid.New(),
		Now:      time.Now(),
	})
	require.NoError(t, err)
	require.Empty(t, ids)

	// Correct worker_id should return the chat's ID.
	ids, err = db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{
		IDs:      []uuid.UUID{chat.ID},
		WorkerID: workerID,
		Now:      time.Now(),
	})
	require.NoError(t, err)
	require.Len(t, ids, 1)
	require.Equal(t, chat.ID, ids[0])
}

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