Function TestChatCostUsers

coderd/exp_chats_test.go:10285–10373

(t *testing.T)

}

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

	seedCtx := testutil.Context(t, testutil.WaitLong)
	client, db := newChatClientWithDatabase(t)
	firstUser := coderdtest.CreateFirstUser(t, client.Client)
	memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
	memberClient := codersdk.NewExperimentalClient(memberClientRaw)
	firstUserRecord, err := db.GetUserByID(dbauthz.AsSystemRestricted(seedCtx), firstUser.UserID)
	require.NoError(t, err)
	modelConfig := createChatModelConfig(t, client)

	adminChat := dbgen.Chat(t, db, database.Chat{
		OrganizationID:    firstUser.OrganizationID,
		OwnerID:           firstUser.UserID,
		LastModelConfigID: modelConfig.ID,
		Title:             "admin chat",
	})
	_ = dbgen.ChatMessage(t, db, database.ChatMessage{
		ChatID:          adminChat.ID,
		ModelConfigID:   uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
		Role:            database.ChatMessageRoleAssistant,
		InputTokens:     sql.NullInt64{Int64: 100, Valid: true},
		OutputTokens:    sql.NullInt64{Int64: 50, Valid: true},
		TotalCostMicros: sql.NullInt64{Int64: 300, Valid: true},
	})

	memberChat := dbgen.Chat(t, db, database.Chat{
		OrganizationID:    firstUser.OrganizationID,
		OwnerID:           member.ID,
		LastModelConfigID: modelConfig.ID,
		Title:             "member chat",
	})
	_ = dbgen.ChatMessage(t, db, database.ChatMessage{
		ChatID:          memberChat.ID,
		ModelConfigID:   uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
		Role:            database.ChatMessageRoleAssistant,
		InputTokens:     sql.NullInt64{Int64: 200, Valid: true},
		OutputTokens:    sql.NullInt64{Int64: 100, Valid: true},
		TotalCostMicros: sql.NullInt64{Int64: 800, Valid: true},
	})

	t.Run("AdminCanListUsers", func(t *testing.T) {
		t.Parallel()

		ctx := testutil.Context(t, testutil.WaitLong)
		resp, err := client.GetChatCostUsers(ctx, codersdk.ChatCostUsersOptions{})
		require.NoError(t, err)
		require.Equal(t, int64(2), resp.Count)
		require.Len(t, resp.Users, 2)
		require.Equal(t, member.ID, resp.Users[0].UserID)
		require.Equal(t, member.Username, resp.Users[0].Username)
		require.Equal(t, int64(800), resp.Users[0].TotalCostMicros)
		require.Equal(t, int64(1), resp.Users[0].MessageCount)
		require.Equal(t, int64(1), resp.Users[0].ChatCount)
		require.Equal(t, firstUser.UserID, resp.Users[1].UserID)
		require.Equal(t, firstUserRecord.Username, resp.Users[1].Username)
		require.Equal(t, int64(300), resp.Users[1].TotalCostMicros)