Atlas/internal/provider/models_test.go

package provider

import "testing"

func TestModelsPath(t *testing.T) {
	openai := ModelsPath(&Config{Protocol: "openai", ServiceType: "chat"})
	if openai != "/v1/models?protocol=openai&service_type=chat" {
		t.Fatalf("openai path: %s", openai)
	}
	replicate := ModelsPath(&Config{Protocol: "replicate", ServiceType: "image_generation"})
	if replicate != "/v1/models?protocol=replicate&service_type=image_generation" {
		t.Fatalf("replicate path: %s", replicate)
	}
}

func TestParseReplicateModels(t *testing.T) {
	body := []byte(`{
		"results": [
			{"owner":"stability-ai","name":"sdxl","description":"SDXL"},
			{"owner":"black-forest-labs","name":"flux-schnell"}
		]
	}`)
	models := parseReplicateModels(body)
	if len(models) != 2 {
		t.Fatalf("expected 2 models, got %d", len(models))
	}
	if models[0].ID != "stability-ai/sdxl" || models[0].Name != "SDXL" {
		t.Fatalf("unexpected first model: %+v", models[0])
	}
	if models[1].ID != "black-forest-labs/flux-schnell" {
		t.Fatalf("unexpected second model id: %s", models[1].ID)
	}
}

func TestParseModelsForProtocol(t *testing.T) {
	openaiBody := []byte(`{"data":[{"id":"gpt-4o","name":"GPT-4o"}]}`)
	if models := parseModelsForProtocol("openai", openaiBody); len(models) != 1 || models[0].ID != "gpt-4o" {
		t.Fatalf("openai parse failed: %+v", models)
	}
	replicateBody := []byte(`{"results":[{"owner":"a","name":"b"}]}`)
	if models := parseModelsForProtocol("replicate", replicateBody); len(models) != 1 || models[0].ID != "a/b" {
		t.Fatalf("replicate parse failed: %+v", models)
	}
}