package openai import ( "bytes" "errors" "io" "net/http" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestConvertOpenAIImageURLsToBase64Disabled(t *testing.T) { t.Setenv("CODEX_IMAGE_URL_TO_B64", "false") body := []byte(`{"created":1,"data":[{"url":"https://example.com/image.png"}]}`) converted, stats, err := convertOpenAIImageURLsToBase64(body, func(string, ...string) (*http.Response, error) { t.Fatal("download must not run while compatibility is disabled") return nil, nil }) require.NoError(t, err) assert.Equal(t, body, converted) assert.Zero(t, stats.Converted) } func TestConvertOpenAIImageURLsToBase64PreservesResponseFields(t *testing.T) { t.Setenv("CODEX_IMAGE_URL_TO_B64", "true") image := []byte("test-image") body := []byte(`{"created":1,"data":[{"url":"https://example.com/image.png","revised_prompt":"kept"}],"usage":{"total_tokens":7},"custom":{"kept":true}}`) converted, stats, err := convertOpenAIImageURLsToBase64(body, func(rawURL string, reason ...string) (*http.Response, error) { assert.Equal(t, "https://example.com/image.png", rawURL) assert.Equal(t, []string{"OpenAI image URL compatibility"}, reason) return &http.Response{ StatusCode: http.StatusOK, ContentLength: int64(len(image)), Body: io.NopCloser(bytes.NewReader(image)), }, nil }) require.NoError(t, err) assert.Equal(t, 1, stats.Converted) assert.Equal(t, int64(len(image)), stats.DownloadedBytes) var payload map[string]any require.NoError(t, common.Unmarshal(converted, &payload)) data := payload["data"].([]any) item := data[0].(map[string]any) assert.Equal(t, "dGVzdC1pbWFnZQ==", item["b64_json"]) assert.Equal(t, "kept", item["revised_prompt"]) assert.NotContains(t, item, "url") assert.Equal(t, float64(7), payload["usage"].(map[string]any)["total_tokens"]) assert.Equal(t, true, payload["custom"].(map[string]any)["kept"]) assert.Contains(t, stats.ServerTiming(), "image-url-compat") } func TestReadImageURLCompatResponseRejectsOversize(t *testing.T) { response := &http.Response{ StatusCode: http.StatusOK, ContentLength: maxImageURLCompatBytes + 1, Body: io.NopCloser(bytes.NewReader([]byte("unused"))), } _, err := readImageURLCompatResponse(response) require.Error(t, err) assert.Contains(t, err.Error(), "exceeds") } func TestDownloadOpenAIImageURLCompatRejectsPrivateLiteral(t *testing.T) { _, err := downloadOpenAIImageURLCompat("https://127.0.0.1/image.png") require.Error(t, err) assert.Contains(t, err.Error(), "private IP address not allowed") } func TestRetryImageURLCompatDownloadRecoversFromTransientError(t *testing.T) { attempts := 0 waits := 0 response, err := retryImageURLCompatDownload(func() (*http.Response, error) { attempts++ if attempts == 1 { return nil, errors.New("EOF") } return &http.Response{ StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewBufferString("image")), }, nil }, func(time.Duration) { waits++ }) require.NoError(t, err) require.NotNil(t, response) assert.Equal(t, 2, attempts) assert.Equal(t, 1, waits) } func TestRetryImageURLCompatDownloadStopsAfterLimit(t *testing.T) { attempts := 0 response, err := retryImageURLCompatDownload(func() (*http.Response, error) { attempts++ return nil, errors.New("EOF") }, func(time.Duration) {}) require.Error(t, err) assert.Nil(t, response) assert.Equal(t, maxImageURLCompatDownloadAttempts, attempts) } func TestRetryImageURLCompatDownloadRejectsEmptyResponse(t *testing.T) { attempts := 0 response, err := retryImageURLCompatDownload(func() (*http.Response, error) { attempts++ return nil, nil }, func(time.Duration) {}) require.Error(t, err) assert.Nil(t, response) assert.Contains(t, err.Error(), "empty response") assert.Equal(t, maxImageURLCompatDownloadAttempts, attempts) }