image_url_compat_test.go

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)
}