ez-api/test/integration_test.go

//go:build integration

package integration

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

type providerResp struct {
	ID uint `json:"ID"`
}

type keyResp struct {
	ID    uint   `json:"ID"`
	Group string `json:"group"`
}

type modelResp struct {
	ID   uint   `json:"ID"`
	Name string `json:"name"`
}

func TestEndToEnd(t *testing.T) {
	apiBase := getenv("E2E_EZAPI_URL", "http://localhost:8080")
	balancerBase := getenv("E2E_BALANCER_URL", "http://localhost:8081")

	client := &http.Client{Timeout: 5 * time.Second}

	// 1) create provider pointing to mock upstream inside compose network
	prov := map[string]string{
		"name":     "mock",
		"type":     "mock",
		"base_url": "http://mock-upstream:8082",
		"api_key":  "mock-upstream-key",
		"group":    "default",
	}
	_ = postJSON(t, client, apiBase+"/providers", prov, new(providerResp))

	// 2) create model metadata
	modelPayload := map[string]interface{}{
		"name":           "mock-model",
		"context_window": 2048,
		"cost_per_token": 0.0,
	}
	_ = postJSON(t, client, apiBase+"/models", modelPayload, new(modelResp))

	// 3) create key bound to provider
	keyPayload := map[string]interface{}{
		"group":      "default",
		"key_secret": "sk-integration",
		"status":     "active",
		"weight":     10,
		"balance":    100,
	}
	postJSON(t, client, apiBase+"/keys", keyPayload, new(keyResp))

	// 4) wait for balancer to refresh snapshot
	waitFor(t, 15*time.Second, func() error {
		models := fetchModels(t, client, balancerBase)
		if len(models) == 0 {
			return fmt.Errorf("no models yet")
		}
		found := false
		for _, m := range models {
			if m == "mock-model" {
				found = true
			}
		}
		if !found {
			return fmt.Errorf("mock-model not visible")
		}
		return nil
	})

	// 5) call chat completions through balancer
	body := map[string]interface{}{
		"model": "mock-model",
		"messages": []map[string]string{
			{"role": "user", "content": "hi"},
		},
	}
	reqBody, _ := json.Marshal(body)
	req, err := http.NewRequest(http.MethodPost, balancerBase+"/v1/chat/completions", bytes.NewReader(reqBody))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+keyPayload["key_secret"].(string))
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	require.Equal(t, http.StatusOK, resp.StatusCode)

	var respBody map[string]interface{}
	data, _ := io.ReadAll(resp.Body)
	require.NoError(t, json.Unmarshal(data, &respBody))
	require.Equal(t, "chat.completion", respBody["object"])
}

func fetchModels(t *testing.T, client *http.Client, balancerBase string) []string {
	req, err := http.NewRequest(http.MethodGet, balancerBase+"/v1/models", nil)
	require.NoError(t, err)
	// No auth required? our balancer requires auth; use test token
	req.Header.Set("Authorization", "Bearer sk-integration")

	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	require.Equal(t, http.StatusOK, resp.StatusCode)

	var payload struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	data, _ := io.ReadAll(resp.Body)
	require.NoError(t, json.Unmarshal(data, &payload))

	out := make([]string, 0, len(payload.Data))
	for _, m := range payload.Data {
		out = append(out, m.ID)
	}
	return out
}

func waitFor(t *testing.T, timeout time.Duration, fn func() error) {
	deadline := time.Now().Add(timeout)
	for {
		if err := fn(); err == nil {
			return
		}
		if time.Now().After(deadline) {
			require.NoError(t, fn())
		}
		time.Sleep(500 * time.Millisecond)
	}
}

func postJSON[T any](t *testing.T, client *http.Client, url string, body interface{}, out *T) *T {
	b, err := json.Marshal(body)
	require.NoError(t, err)

	req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(b))
	require.NoError(t, err)
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	require.Equal(t, http.StatusCreated, resp.StatusCode)

	data, _ := io.ReadAll(resp.Body)
	if len(data) > 0 {
		require.NoError(t, json.Unmarshal(data, out))
	}
	return out
}

func getenv(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}