Function TestUserAIProviderKeys

coderd/exp_chats_test.go:2081–2259

(t *testing.T)

}

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

	createOpenAIProvider := func(t *testing.T, client *codersdk.ExperimentalClient, name string, enabled bool, apiKeys ...string) codersdk.AIProvider {
		t.Helper()

		provider, err := client.CreateAIProvider(testutil.Context(t, testutil.WaitLong), codersdk.CreateAIProviderRequest{
			Type:    codersdk.AIProviderTypeOpenAI,
			Name:    name,
			Enabled: enabled,
			BaseURL: "https://api.openai.example.com/v1",
			APIKeys: apiKeys,
		})
		require.NoError(t, err)
		return provider
	}

	findUserAIProviderKeyConfig := func(
		t *testing.T,
		configs []codersdk.UserAIProviderKeyConfig,
		providerID uuid.UUID,
	) *codersdk.UserAIProviderKeyConfig {
		t.Helper()

		for i := range configs {
			if configs[i].Provider.ID == providerID {
				return &configs[i]
			}
		}
		return nil
	}

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

		ctx := testutil.Context(t, testutil.WaitLong)
		adminClient := newChatClient(t)
		firstUser := coderdtest.CreateFirstUser(t, adminClient.Client)
		memberClientRaw, _ := coderdtest.CreateAnotherUser(t, adminClient.Client, firstUser.OrganizationID)
		memberClient := codersdk.NewExperimentalClient(memberClientRaw)

		provider := createOpenAIProvider(t, adminClient, "test-user-key-"+uuid.NewString(), true, "test-provider-api-key")

		configs, err := memberClient.ListUserAIProviderKeyConfigs(ctx, "me")
		require.NoError(t, err)
		cfg := findUserAIProviderKeyConfig(t, configs, provider.ID)
		require.NotNil(t, cfg)
		require.False(t, cfg.HasUserAPIKey)
		require.True(t, cfg.HasProviderAPIKey)
		require.True(t, cfg.BYOKEnabled)

		cfgValue, err := memberClient.UpsertUserAIProviderKey(ctx, "me", provider.ID, codersdk.CreateUserAIProviderKeyRequest{APIKey: "test-user-api-key"})
		require.NoError(t, err)
		require.Equal(t, provider.ID, cfgValue.Provider.ID)
		require.True(t, cfgValue.HasUserAPIKey)
		require.True(t, cfgValue.HasProviderAPIKey)
		require.True(t, cfgValue.BYOKEnabled)