llm-intelligence/scripts/official_pricing_import_common_test.go

//go:build llm_script

package main

import (
	"net/http"
	"net/http/httptest"
	"os"
	"sync/atomic"
	"testing"
	"time"
)

func TestFetchRawPricingPageRetriesTransientStatus(t *testing.T) {
	var attempts int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		current := atomic.AddInt32(&attempts, 1)
		if current == 1 {
			http.Error(w, "temporary", http.StatusServiceUnavailable)
			return
		}
		_, _ = w.Write([]byte("ok"))
	}))
	defer server.Close()

	client := &http.Client{Timeout: 2 * time.Second}
	body, err := fetchRawPricingPage(server.URL, "", client)
	if err != nil {
		t.Fatalf("fetchRawPricingPage returned error: %v", err)
	}
	if body != "ok" {
		t.Fatalf("body = %q, want ok", body)
	}
	if got := atomic.LoadInt32(&attempts); got != 2 {
		t.Fatalf("attempts = %d, want 2", got)
	}
}

func TestIsRetriablePricingFetchErrorRecognizesEOF(t *testing.T) {
	if !isRetriablePricingFetchError(errString("unexpected EOF")) {
		t.Fatalf("expected EOF to be retriable")
	}
	if isRetriablePricingFetchError(errString("bad request")) {
		t.Fatalf("expected bad request to be non-retriable")
	}
}

func TestFetchRawPricingPageFallsBackWithoutProxyOnRetriableProxyFailure(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("ok"))
	}))
	defer server.Close()

	oldHTTPProxy, hadHTTPProxy := os.LookupEnv("HTTP_PROXY")
	oldHTTPSProxy, hadHTTPSProxy := os.LookupEnv("HTTPS_PROXY")
	oldNoProxy, hadNoProxy := os.LookupEnv("NO_PROXY")
	defer func() {
		if hadHTTPProxy {
			_ = os.Setenv("HTTP_PROXY", oldHTTPProxy)
		} else {
			_ = os.Unsetenv("HTTP_PROXY")
		}
		if hadHTTPSProxy {
			_ = os.Setenv("HTTPS_PROXY", oldHTTPSProxy)
		} else {
			_ = os.Unsetenv("HTTPS_PROXY")
		}
		if hadNoProxy {
			_ = os.Setenv("NO_PROXY", oldNoProxy)
		} else {
			_ = os.Unsetenv("NO_PROXY")
		}
	}()

	_ = os.Setenv("HTTP_PROXY", "http://127.0.0.1:1")
	_ = os.Unsetenv("HTTPS_PROXY")
	_ = os.Unsetenv("NO_PROXY")

	client := &http.Client{Timeout: 2 * time.Second}
	body, err := fetchRawPricingPage(server.URL, "", client)
	if err != nil {
		t.Fatalf("fetchRawPricingPage returned error with proxy fallback enabled: %v", err)
	}
	if body != "ok" {
		t.Fatalf("body = %q, want ok", body)
	}
}

func TestFallbackModalityCanonicalizesAliases(t *testing.T) {
	if got := fallbackModality("image"); got != "vision" {
		t.Fatalf("fallbackModality(image) = %q, want vision", got)
	}
	if got := fallbackModality(" "); got != "text" {
		t.Fatalf("fallbackModality(blank) = %q, want text", got)
	}
}

type errString string

func (e errString) Error() string { return string(e) }