iop/packages/go/config/usage_attribution_config_test.go

package config_test

import (
	"os"
	"path/filepath"
	"strings"
	"testing"

	"iop/packages/go/config"
)

func loadUsageAttributionConfig(t *testing.T, content string) (*config.EdgeConfig, error) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "edge.yaml")
	if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
		t.Fatalf("write config: %v", err)
	}
	return config.LoadEdge(path)
}

func TestModelUsageAttributionValidationAndDefault(t *testing.T) {
	const basePrefix = `
models:
  - id: "logical-model"
`
	const baseSuffix = `
    providers:
      provider-a: "served-model"
nodes:
  - id: "node-a"
    providers:
      - id: "provider-a"
        type: "vllm"
        category: "api"
        models:
          - "served-model"
`

	for _, tc := range []struct {
		name        string
		policyYAML  string
		want        string
		wantErrPart string
	}{
		{name: "omitted defaults to provider", want: config.UsageAttributionProvider},
		{name: "provider is accepted", policyYAML: `    usage_attribution: "provider"` + "\n", want: config.UsageAttributionProvider},
		{name: "model group is accepted", policyYAML: `    usage_attribution: "model_group"` + "\n", want: config.UsageAttributionModelGroup},
		{name: "unknown policy is rejected", policyYAML: `    usage_attribution: "request_model"` + "\n", wantErrPart: "usage_attribution"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			cfg, err := loadUsageAttributionConfig(t, basePrefix+tc.policyYAML+baseSuffix)
			if tc.wantErrPart != "" {
				if err == nil {
					t.Fatal("expected config load error")
				}
				if !strings.Contains(err.Error(), tc.wantErrPart) {
					t.Fatalf("error %q does not contain %q", err, tc.wantErrPart)
				}
				return
			}
			if err != nil {
				t.Fatalf("load config: %v", err)
			}
			if got := cfg.Models[0].EffectiveUsageAttribution(); got != tc.want {
				t.Fatalf("effective attribution = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestOpenAIDirectProviderBindingValidation(t *testing.T) {
	for _, tc := range []struct {
		name        string
		openAIYAML  string
		wantErrPart string
		check       func(*testing.T, *config.EdgeConfig)
	}{
		{
			name: "disabled surface does not require binding",
			openAIYAML: `
openai:
  enabled: false
`,
		},
		{
			name: "top level fallback binding",
			openAIYAML: `
openai:
  enabled: true
  provider_id: "provider-fallback"
  adapter: "cli"
  target: "codex"
`,
			check: func(t *testing.T, cfg *config.EdgeConfig) {
				if cfg.OpenAI.ProviderID != "provider-fallback" {
					t.Fatalf("provider_id = %q", cfg.OpenAI.ProviderID)
				}
			},
		},
		{
			name: "route override and fallback binding",
			openAIYAML: `
openai:
  enabled: true
  provider_id: "provider-fallback"
  model_routes:
    - model: "special"
      provider_id: "provider-route"
      adapter: "cli"
      target: "codex"
    - model: "inherited"
      provider_id: "  "
      adapter: "cli"
      target: "codex-exec"
`,
			check: func(t *testing.T, cfg *config.EdgeConfig) {
				if got := cfg.OpenAI.ModelRoutes[0].ProviderID; got != "provider-route" {
					t.Fatalf("route provider_id = %q", got)
				}
				if got := strings.TrimSpace(cfg.OpenAI.ModelRoutes[1].ProviderID); got != "" {
					t.Fatalf("blank route provider_id = %q", got)
				}
			},
		},
		{
			name: "missing fallback binding is rejected",
			openAIYAML: `
openai:
  enabled: true
  adapter: "cli"
  target: "codex"
`,
			wantErrPart: "openai.provider_id",
		},
		{
			name: "blank fallback binding is rejected",
			openAIYAML: `
openai:
  enabled: true
  provider_id: "  "
  model_routes:
    - model: "special"
      provider_id: "provider-route"
      adapter: "cli"
      target: "codex"
`,
			wantErrPart: "openai.provider_id",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			cfg, err := loadUsageAttributionConfig(t, tc.openAIYAML)
			if tc.wantErrPart != "" {
				if err == nil {
					t.Fatal("expected config load error")
				}
				if !strings.Contains(err.Error(), tc.wantErrPart) {
					t.Fatalf("error %q does not contain %q", err, tc.wantErrPart)
				}
				return
			}
			if err != nil {
				t.Fatalf("load config: %v", err)
			}
			if tc.check != nil {
				tc.check(t, cfg)
			}
		})
	}
}