iop/packages/go/agentprovider/catalog/discovery_test.go

package catalog

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"iop/packages/go/agentconfig"
)

func TestDiscoveryStateTableWithIsolatedPATH(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell fixture requires Unix")
	}
	binDir := t.TempDir()
	writeExecutable(t, filepath.Join(binDir, "ready-cli"), fakeProbeScript("ready"))
	writeExecutable(t, filepath.Join(binDir, "unauth-cli"), fakeProbeScript("unauthenticated"))
	writeExecutable(t, filepath.Join(binDir, "unsupported-cli"), fakeProbeScript("unsupported"))
	writeExecutable(t, filepath.Join(binDir, "error-cli"), fakeProbeScript("error"))
	t.Setenv("PATH", binDir)

	cfg := stateTableCatalog()
	discoverer, err := NewDiscoverer(cfg, nil)
	if err != nil {
		t.Fatalf("NewDiscoverer: %v", err)
	}
	results := discoverer.Discover(context.Background())
	if got, want := len(results), 5; got != want {
		t.Fatalf("result count = %d, want %d", got, want)
	}

	wantStates := map[string]ReadinessState{
		"error-profile":       StateProbeError,
		"missing-profile":     StateMissingBinary,
		"ready-profile":       StateReady,
		"unauth-profile":      StateUnauthenticated,
		"unsupported-profile": StateUnsupportedModel,
	}
	for _, result := range results {
		if want := wantStates[result.ProfileID]; result.State != want {
			t.Errorf("%s state = %q, want %q; detail=%q", result.ProfileID, result.State, want, result.Detail)
		}
		switch result.State {
		case StateMissingBinary:
			if !errors.Is(result.Error, ErrBinaryMissing) {
				t.Errorf("%s error = %v, want ErrBinaryMissing", result.ProfileID, result.Error)
			}
		case StateUnauthenticated:
			if !errors.Is(result.Error, ErrAuthenticationRequired) {
				t.Errorf("%s error = %v, want ErrAuthenticationRequired", result.ProfileID, result.Error)
			}
		case StateUnsupportedModel:
			if !errors.Is(result.Error, ErrModelUnsupported) {
				t.Errorf("%s error = %v, want ErrModelUnsupported", result.ProfileID, result.Error)
			}
		case StateProbeError:
			if !errors.Is(result.Error, ErrProbeFailed) {
				t.Errorf("%s error = %v, want ErrProbeFailed", result.ProfileID, result.Error)
			}
			if strings.Contains(result.Detail, "probe-secret") {
				t.Errorf("%s leaked probe secret: %q", result.ProfileID, result.Detail)
			}
		}
	}
	for index := 1; index < len(results); index++ {
		if results[index-1].ProfileID > results[index].ProfileID {
			t.Fatalf("results not sorted: %q before %q", results[index-1].ProfileID, results[index].ProfileID)
		}
	}
}

func TestDiscoveryTimeoutAndCancellationAreProbeErrors(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell fixture requires Unix")
	}
	binDir := t.TempDir()
	writeExecutable(t, filepath.Join(binDir, "slow-cli"), `#!/bin/sh
case "$1" in
  --version) echo "slow 1.0" ;;
  auth) exec /bin/sleep 5 ;;
  models) echo model-a ;;
esac
`)
	t.Setenv("PATH", binDir)

	cfg := oneProfileCatalog("slow", "slow-cli", 25)
	discoverer, err := NewDiscoverer(cfg, nil)
	if err != nil {
		t.Fatalf("NewDiscoverer: %v", err)
	}
	result, err := discoverer.DiscoverProfile(context.Background(), "slow-profile")
	if result.State != StateProbeError || !errors.Is(err, ErrProbeFailed) {
		t.Fatalf("timeout state/error = %q/%v", result.State, err)
	}
	if !strings.Contains(result.Detail, "context deadline exceeded") {
		t.Fatalf("timeout detail = %q", result.Detail)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	result, err = discoverer.DiscoverProfile(ctx, "slow-profile")
	if result.State != StateProbeError || !errors.Is(err, ErrProbeFailed) {
		t.Fatalf("cancel state/error = %q/%v", result.State, err)
	}
	if !strings.Contains(result.Detail, "context canceled") {
		t.Fatalf("cancel detail = %q", result.Detail)
	}
}

func stateTableCatalog() agentconfig.Catalog {
	provider := func(id, command string) agentconfig.Provider {
		return agentconfig.Provider{
			ID:           id,
			Command:      command,
			VersionProbe: agentconfig.CommandProbe{Args: []string{"--version"}, TimeoutMS: 500},
			Authentication: agentconfig.AuthenticationProbe{
				Args:                   []string{"auth"},
				TimeoutMS:              500,
				SuccessPattern:         `^authenticated$`,
				UnauthenticatedPattern: `(?i)not logged in`,
			},
			ModelProbe:   agentconfig.ModelProbe{Args: []string{"models"}, TimeoutMS: 500},
			Capabilities: []string{"run", "status"},
		}
	}
	providers := []agentconfig.Provider{
		provider("error", "error-cli"),
		provider("missing", "missing-cli"),
		provider("ready", "ready-cli"),
		provider("unauth", "unauth-cli"),
		provider("unsupported", "unsupported-cli"),
	}
	models := make([]agentconfig.Model, 0, len(providers))
	profiles := make([]agentconfig.Profile, 0, len(providers))
	for _, item := range providers {
		models = append(models, agentconfig.Model{
			ID: item.ID + "-model", Provider: item.ID, Target: "model-a",
		})
		profiles = append(profiles, agentconfig.Profile{
			ID: item.ID + "-profile", Provider: item.ID, Model: item.ID + "-model",
			Capabilities: []string{"run", "status"},
		})
	}
	return agentconfig.Catalog{
		Version:   agentconfig.SchemaVersion,
		Providers: providers,
		Models:    models,
		Profiles:  profiles,
	}
}

func oneProfileCatalog(id, command string, authTimeoutMS int) agentconfig.Catalog {
	return agentconfig.Catalog{
		Version: agentconfig.SchemaVersion,
		Providers: []agentconfig.Provider{{
			ID:           id,
			Command:      command,
			VersionProbe: agentconfig.CommandProbe{Args: []string{"--version"}, TimeoutMS: 500},
			Authentication: agentconfig.AuthenticationProbe{
				Args:      []string{"auth"},
				TimeoutMS: authTimeoutMS,
			},
			ModelProbe:   agentconfig.ModelProbe{Args: []string{"models"}, TimeoutMS: 500},
			Capabilities: []string{"run", "status"},
		}},
		Models: []agentconfig.Model{{
			ID: id + "-model", Provider: id, Target: "model-a",
		}},
		Profiles: []agentconfig.Profile{{
			ID: id + "-profile", Provider: id, Model: id + "-model",
			Capabilities: []string{"run", "status"},
		}},
	}
}

func fakeProbeScript(state string) string {
	switch state {
	case "ready":
		return `#!/bin/sh
case "$1" in
  --version) echo "fake 1.0" ;;
  auth) echo authenticated ;;
  models) echo model-a ;;
esac
`
	case "unauthenticated":
		return `#!/bin/sh
case "$1" in
  --version) echo "fake 1.0" ;;
  auth) echo "not logged in"; exit 0 ;;
  models) echo model-a ;;
esac
`
	case "unsupported":
		return `#!/bin/sh
case "$1" in
  --version) echo "fake 1.0" ;;
  auth) echo authenticated ;;
  models) echo model-b ;;
esac
`
	default:
		return `#!/bin/sh
case "$1" in
  --version) echo "fake 1.0" ;;
  auth) echo "api_key=probe-secret backend unavailable"; exit 2 ;;
  models) echo model-a ;;
esac
`
	}
}

func writeExecutable(t *testing.T, path, content string) {
	t.Helper()
	if err := os.WriteFile(path, []byte(content), 0o755); err != nil {
		t.Fatalf("write executable: %v", err)
	}
}

func TestDiscoverProfileUnknown(t *testing.T) {
	discoverer, err := NewDiscoverer(oneProfileCatalog("ready", "ready-cli", 100), &fakeRunner{})
	if err != nil {
		t.Fatalf("NewDiscoverer: %v", err)
	}
	if _, err := discoverer.DiscoverProfile(context.Background(), "missing"); err == nil {
		t.Fatal("DiscoverProfile unexpectedly succeeded")
	}
}

type fakeRunner struct{}

func (*fakeRunner) LookPath(command string) (string, error) {
	return command, nil
}

func (*fakeRunner) Run(_ context.Context, _ string, _ []string) (RunResult, error) {
	return RunResult{Output: "ok"}, nil
}

func TestDiagnosticTruncatesAfterRedaction(t *testing.T) {
	value := "api_key=probe-secret " + strings.Repeat("x", 4096)
	got := diagnostic(value)
	if strings.Contains(got, "probe-secret") {
		t.Fatalf("diagnostic leaked secret: %q", got)
	}
	if len(got) > 2051 {
		t.Fatalf("diagnostic length = %d", len(got))
	}
}

func TestRunProbeHonorsParentDeadline(t *testing.T) {
	discoverer := &Discoverer{runner: blockingRunner{}}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	_, err := discoverer.runProbe(ctx, "blocked", []string{"auth"}, 10_000)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("runProbe error = %v", err)
	}
}

type blockingRunner struct{}

func (blockingRunner) LookPath(command string) (string, error) {
	return command, nil
}

func (blockingRunner) Run(ctx context.Context, _ string, _ []string) (RunResult, error) {
	<-ctx.Done()
	return RunResult{}, ctx.Err()
}