easyai-ai-gateway/apps/api/internal/runner/proxy_test.go

package runner

import (
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestProviderHTTPClientTimeoutAllowsLongRunningMediaRequests(t *testing.T) {
	client := newHTTPClient(nil)
	if client.Timeout != 10*time.Minute {
		t.Fatalf("unexpected provider HTTP timeout: got %s want %s", client.Timeout, 10*time.Minute)
	}
	for _, test := range []struct {
		kind string
		want time.Duration
	}{
		{kind: "images.generations", want: 20 * time.Minute},
		{kind: "images.edits", want: 20 * time.Minute},
		{kind: "videos.generations", want: 30 * time.Minute},
		{kind: "chat.completions", want: 10 * time.Minute},
	} {
		if got := providerHTTPClientForKind(client, test.kind).Timeout; got != test.want {
			t.Fatalf("provider HTTP timeout for %s: got %s want %s", test.kind, got, test.want)
		}
	}
	transport, ok := client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("provider transport type = %T, want *http.Transport", client.Transport)
	}
	if transport.MaxIdleConns != providerHTTPMaxIdleConnections ||
		transport.MaxIdleConnsPerHost != providerHTTPMaxIdleConnectionsPerHost {
		t.Fatalf(
			"provider idle connection pool = %d/%d, want %d/%d",
			transport.MaxIdleConns,
			transport.MaxIdleConnsPerHost,
			providerHTTPMaxIdleConnections,
			providerHTTPMaxIdleConnectionsPerHost,
		)
	}
}

func TestPlatformProxyModeNoneIgnoresEnvironmentProxy(t *testing.T) {
	var proxyHits int
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyHits++
		w.WriteHeader(http.StatusTeapot)
	}))
	defer proxy.Close()
	t.Setenv("HTTP_PROXY", proxy.URL)
	t.Setenv("http_proxy", proxy.URL)

	var targetHits int
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		targetHits++
		_, _ = w.Write([]byte("ok"))
	}))
	defer target.Close()

	client, err := testProxyService(config.Config{}).httpClientForCandidate(store.RuntimeModelCandidate{
		PlatformConfig: map[string]any{"networkProxy": map[string]any{"mode": "none"}},
	}, false, "chat.completions")
	if err != nil {
		t.Fatalf("build http client: %v", err)
	}
	resp, err := client.Get(target.URL)
	if err != nil {
		t.Fatalf("get target: %v", err)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK || targetHits != 1 || proxyHits != 0 {
		t.Fatalf("unexpected status=%d targetHits=%d proxyHits=%d", resp.StatusCode, targetHits, proxyHits)
	}
}

func TestPlatformProxyModeCustomUsesConfiguredHTTPProxy(t *testing.T) {
	var targetHits int
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		targetHits++
		_, _ = w.Write([]byte("target"))
	}))
	defer target.Close()

	var proxyHits int
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyHits++
		if r.URL.String() != target.URL && r.URL.String() != target.URL+"/" {
			t.Fatalf("proxy received unexpected target URL %q", r.URL.String())
		}
		_, _ = w.Write([]byte("proxied"))
	}))
	defer proxy.Close()

	client, err := testProxyService(config.Config{}).httpClientForCandidate(store.RuntimeModelCandidate{
		PlatformConfig: map[string]any{"networkProxy": map[string]any{"mode": "custom", "httpProxy": proxy.URL}},
	}, false, "chat.completions")
	if err != nil {
		t.Fatalf("build http client: %v", err)
	}
	resp, err := client.Get(target.URL)
	if err != nil {
		t.Fatalf("get target through proxy: %v", err)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK || proxyHits != 1 || targetHits != 0 {
		t.Fatalf("unexpected status=%d proxyHits=%d targetHits=%d", resp.StatusCode, proxyHits, targetHits)
	}
}

func TestPlatformProxyBypassIDUsesDirectConnection(t *testing.T) {
	var targetHits int
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		targetHits++
		_, _ = w.Write([]byte("target"))
	}))
	defer target.Close()

	var proxyHits int
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyHits++
		w.WriteHeader(http.StatusProxyAuthRequired)
	}))
	defer proxy.Close()

	client, err := testProxyService(config.Config{
		PlatformProxyBypassIDs: "other-platform, official-gemini-platform ",
	}).httpClientForCandidate(store.RuntimeModelCandidate{
		PlatformID:     "official-gemini-platform",
		PlatformConfig: map[string]any{"networkProxy": map[string]any{"mode": "custom", "httpProxy": proxy.URL}},
	}, false, "chat.completions")
	if err != nil {
		t.Fatalf("build bypassed http client: %v", err)
	}
	resp, err := client.Get(target.URL)
	if err != nil {
		t.Fatalf("get target directly: %v", err)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK || targetHits != 1 || proxyHits != 0 {
		t.Fatalf("unexpected status=%d targetHits=%d proxyHits=%d", resp.StatusCode, targetHits, proxyHits)
	}
}

func TestPlatformProxyModeGlobalUsesConfiguredGlobalHTTPProxy(t *testing.T) {
	var targetHits int
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		targetHits++
		_, _ = w.Write([]byte("target"))
	}))
	defer target.Close()

	var proxyHits int
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyHits++
		_, _ = w.Write([]byte("proxied"))
	}))
	defer proxy.Close()

	client, err := testProxyService(config.Config{GlobalHTTPProxy: proxy.URL}).httpClientForCandidate(store.RuntimeModelCandidate{
		PlatformConfig: map[string]any{"networkProxy": map[string]any{"mode": "global"}},
	}, false, "chat.completions")
	if err != nil {
		t.Fatalf("build http client: %v", err)
	}
	resp, err := client.Get(target.URL)
	if err != nil {
		t.Fatalf("get target through global proxy: %v", err)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK || proxyHits != 1 || targetHits != 0 {
		t.Fatalf("unexpected status=%d proxyHits=%d targetHits=%d", resp.StatusCode, proxyHits, targetHits)
	}
}

func testProxyService(cfg config.Config) *Service {
	return &Service{cfg: cfg, httpClients: newHTTPClientCache()}
}