account_test_models_test.go 3.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. package service
  2. import (
  3. "context"
  4. "net/http"
  5. "testing"
  6. "github.com/stretchr/testify/require"
  7. )
  8. func TestFetchOpenAIAccountModelsOAuthPopulatesPickerFields(t *testing.T) {
  9. _, calls := newCodexModelsOAuthCacheServer(t, `{"models":[{"slug":"new-oauth-model"},{"slug":"gpt-6-astra"}]}`)
  10. gateway := &OpenAIGatewayService{}
  11. svc := &AccountTestService{openaiGatewayService: gateway}
  12. account := newCodexModelsTestAccount()
  13. ctx := context.Background()
  14. before, err := gateway.FetchOpenAIModelsList(ctx, account)
  15. require.NoError(t, err)
  16. models, err := svc.FetchOpenAIAccountModels(ctx, account)
  17. require.NoError(t, err)
  18. require.Greater(t, len(models), 2)
  19. for i, id := range []string{"new-oauth-model", "gpt-6-astra"} {
  20. require.Equal(t, id, models[i].ID)
  21. require.Equal(t, id, models[i].DisplayName)
  22. require.Equal(t, "model", models[i].Type)
  23. }
  24. ids := make([]string, 0, len(models))
  25. for _, model := range models {
  26. ids = append(ids, model.ID)
  27. }
  28. require.Contains(t, ids, "gpt-image-2.5-flare")
  29. require.Contains(t, ids, "gpt-image-2.5-sunburst")
  30. after, err := gateway.FetchOpenAIModelsList(ctx, account)
  31. require.NoError(t, err)
  32. require.Equal(t, before.Body, after.Body, "picker fields must not change the shared catalog")
  33. require.NotContains(t, string(after.Body), "display_name")
  34. require.EqualValues(t, 1, calls.Load(), "picker must reuse the shared discovery cache")
  35. }
  36. func TestFetchOpenAIAccountModelsAPIKeyPopulatesPickerFields(t *testing.T) {
  37. gateway := newCodexModelsAPIKeyTestService(&codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
  38. return ordinaryModelsUpstreamResponse(`{"data":[
  39. {"id":"new-api-model","owned_by":"provider","created":123},
  40. {"id":"blank-label","display_name":" ","type":""},
  41. {"id":"named-model","display_name":"Provider Model","type":"model"}
  42. ]}`), nil
  43. }})
  44. svc := &AccountTestService{openaiGatewayService: gateway}
  45. models, err := svc.FetchOpenAIAccountModels(context.Background(), newCodexModelsAPIKeyTestAccount("https://models.example/v1"))
  46. require.NoError(t, err)
  47. require.Len(t, models, 3)
  48. for i, name := range []string{"new-api-model", "blank-label", "Provider Model"} {
  49. require.Equal(t, name, models[i].DisplayName)
  50. require.Equal(t, "model", models[i].Type)
  51. }
  52. require.Equal(t, "provider", models[0].OwnedBy)
  53. require.EqualValues(t, 123, models[0].Created)
  54. require.Equal(t, "named-model", models[2].ID)
  55. }
  56. func TestFetchOpenAIAccountModelsPreservesEmptyCatalog(t *testing.T) {
  57. gateway := newCodexModelsAPIKeyTestService(&codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
  58. return ordinaryModelsUpstreamResponse(`{"data":[]}`), nil
  59. }})
  60. svc := &AccountTestService{openaiGatewayService: gateway}
  61. models, err := svc.FetchOpenAIAccountModels(context.Background(), newCodexModelsAPIKeyTestAccount("https://models.example/v1"))
  62. require.NoError(t, err)
  63. require.Empty(t, models, "an empty upstream catalog must not become a static model list")
  64. }
  65. func TestFetchOpenAIAccountModelsOAuthRespectsImageAllowlist(t *testing.T) {
  66. newCodexModelsOAuthCacheServer(t, `{"models":[{"slug":"gpt-6-astra"}]}`)
  67. svc := &AccountTestService{openaiGatewayService: &OpenAIGatewayService{}}
  68. account := newCodexModelsTestAccount()
  69. account.Credentials["model_mapping"] = map[string]any{"gpt-image-2.5-flare": "gpt-image-2.5-flare"}
  70. models, err := svc.FetchOpenAIAccountModels(context.Background(), account)
  71. require.NoError(t, err)
  72. ids := []string{}
  73. for _, model := range models {
  74. ids = append(ids, model.ID)
  75. }
  76. require.Contains(t, ids, "gpt-image-2.5-flare")
  77. require.NotContains(t, ids, "gpt-image-2.5-sunburst")
  78. }