Function TestCreateChatRejectsWhenUsageLimitReached

coderd/x/chatd/chatd_test.go:2447–2517

(t *testing.T)

}

func TestCreateChatRejectsWhenUsageLimitReached(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)

	_, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{
		Enabled:            true,
		DefaultLimitMicros: 100,
		Period:             string(codersdk.ChatUsageLimitPeriodDay),
	})
	require.NoError(t, err)

	existingChat := dbgen.Chat(t, db, database.Chat{
		OrganizationID:    org.ID,
		OwnerID:           user.ID,
		Title:             "existing-limit-chat",
		LastModelConfigID: model.ID,
	})

	assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
		codersdk.ChatMessageText("assistant"),
	})
	require.NoError(t, err)

	_ = dbgen.ChatMessage(t, db, database.ChatMessage{
		ChatID:          existingChat.ID,
		ModelConfigID:   uuid.NullUUID{UUID: model.ID, Valid: true},
		Role:            database.ChatMessageRoleAssistant,
		ContentVersion:  chatprompt.CurrentContentVersion,
		Content:         assistantContent,
		TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
	})

	beforeChats, err := db.GetChats(ctx, database.GetChatsParams{
		OwnedOnly: true,
		ViewerID:  user.ID,
		AfterID:   uuid.Nil,
		OffsetOpt: 0,
		LimitOpt:  100,
	})
	require.NoError(t, err)
	require.Len(t, beforeChats, 1)

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

	var limitErr *chatd.UsageLimitExceededError
	require.ErrorAs(t, err, &limitErr)