llm-intelligence/scripts/fetch_multi_source_test.go

//go:build llm_script

package main

import (
	"bytes"
	"testing"
)

type fakeSource struct {
	name   string
	prices []ModelPricing
	err    error
}

func (s fakeSource) Name() string { return s.name }

func (s fakeSource) FetchPricing() ([]ModelPricing, error) { return s.prices, s.err }

func (s fakeSource) SourceType() string { return "official" }

func TestBuildSourcesFiltersRequestedNames(t *testing.T) {
	sources, err := buildSources("", []string{"moonshot", "openai"})
	if err != nil {
		t.Fatalf("buildSources returned error: %v", err)
	}

	if len(sources) != 2 {
		t.Fatalf("expected 2 sources, got %d", len(sources))
	}

	if sources[0].Name() != "Moonshot" || sources[1].Name() != "OpenAI" {
		t.Fatalf("unexpected source order: %s, %s", sources[0].Name(), sources[1].Name())
	}
}

func TestBuildSourcesRejectsUnknownNames(t *testing.T) {
	_, err := buildSources("", []string{"moonshot", "unknown"})
	if err == nil {
		t.Fatal("expected error for unknown source")
	}
}

func TestRunCollectorDryRunSkipsDatabaseWrite(t *testing.T) {
	cfg := runConfig{DryRun: true}
	var out bytes.Buffer
	writeCalled := false

	err := runCollector(
		cfg,
		[]DataSource{
			fakeSource{
				name: "Moonshot",
				prices: []ModelPricing{
					{ModelID: "kimi-k2.6", ProviderCountry: "CN", Currency: "CNY"},
					{ModelID: "kimi-k2-0905-preview", ProviderCountry: "CN", Currency: "CNY"},
				},
			},
			fakeSource{
				name: "OpenAI",
				prices: []ModelPricing{
					{ModelID: "gpt-5.5", ProviderCountry: "US", Currency: "USD"},
				},
			},
		},
		func([]ModelPricing) error {
			writeCalled = true
			return nil
		},
		&out,
	)
	if err != nil {
		t.Fatalf("runCollector returned error: %v", err)
	}

	if writeCalled {
		t.Fatal("expected dry-run to skip database write")
	}

	output := out.String()
	if output == "" {
		t.Fatal("expected dry-run summary output")
	}
	if !bytes.Contains(out.Bytes(), []byte("sources=2")) {
		t.Fatalf("expected sources summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("models=3")) {
		t.Fatalf("expected model summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("currencies=CNY:2,USD:1")) {
		t.Fatalf("expected currency summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("selected_source_keys=moonshot,openai")) {
		t.Fatalf("expected selected source keys in summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("successful_source_keys=moonshot,openai")) {
		t.Fatalf("expected successful source keys in summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("failed_source_keys=none")) {
		t.Fatalf("expected failed source keys in summary, got %q", output)
	}
}

func TestRunCollectorReportsFailedSourceKeys(t *testing.T) {
	cfg := runConfig{DryRun: true}
	var out bytes.Buffer

	err := runCollector(
		cfg,
		[]DataSource{
			fakeSource{
				name: "Moonshot",
				prices: []ModelPricing{
					{ModelID: "kimi-k2.6", ProviderCountry: "CN", Currency: "CNY"},
				},
			},
			fakeSource{
				name: "OpenAI",
				err:  bytes.ErrTooLarge,
			},
		},
		nil,
		&out,
	)
	if err != nil {
		t.Fatalf("runCollector returned error: %v", err)
	}

	output := out.String()
	if !bytes.Contains(out.Bytes(), []byte("successful_source_keys=moonshot")) {
		t.Fatalf("expected successful source keys in summary, got %q", output)
	}
	if !bytes.Contains(out.Bytes(), []byte("failed_source_keys=openai")) {
		t.Fatalf("expected failed source keys in summary, got %q", output)
	}
}

func TestPricingMetadataClassifiesSourceType(t *testing.T) {
	freeTier := pricingMetadata(ModelPricing{OperatorType: "official", IsFree: true})
	if freeTier.SourceType != "free_tier" {
		t.Fatalf("expected free_tier, got %q", freeTier.SourceType)
	}
	if freeTier.FreeQuota == "" {
		t.Fatal("expected free tier quota description")
	}

	reseller := pricingMetadata(ModelPricing{OperatorType: "reseller"})
	if reseller.SourceType != "reseller" {
		t.Fatalf("expected reseller, got %q", reseller.SourceType)
	}
}