iop/scripts/inventory-query/main_test.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"gopkg.in/yaml.v3"
)

// fixtureYAML is a minimal inventory shape covering both dev and dev-corp patterns.
const fixtureYAML = `
inventory_id: inventory-dev
common_inventory: agent-test/inventory.yaml
test_env: dev
profile: dev-runtime-provider-pool
last_updated_at: "2026-07-13"
source:
  remote_runner:
    ssh: toki@toki-labs.com
edge:
  id: edge-toki-labs-dev
build:
  binaries:
    edge: build/dev-runtime/bin/edge
model:
  alias: ornith:35b
  aliases:
    gemma4:26b:
      capacity_total: 5
      providers:
        - id: mac-mlx-vllm
      served_model: mlx-community/gemma-4-26b-a4b-it-nvfp4
    ornith:35b:
      capacity_total: 8
      providers:
        - id: corp-dgx-spark-01-ornith
        - id: corp-dgx-spark-02-ornith
  active_edge_model_group:
    providers:
      gx10-vllm:
        capacity_total: 4
  qwen3_6_reference:
    alias: qwen3.6:35b
    id: qwen3-6-ref
    provider: mac-mlx-vllm
nodes:
  - id: mac-codex-node
    alias: mac-codex
    providers:
      - id: mac-mlx-vllm
        type: vllm-mlx
  - id: gx10-vllm-node
    alias: gx10-vllm
    provider:
      id: gx10-vllm
      type: vllm
  - id: onexplayer-lemonade-node
    alias: onexplayer-lemonade
    direct_providers:
      - id: ornith-direct
        family: spark_ornith
  - id: node-3
    alias: node-3
  - id: node-4
    alias: node-4
  - id: node-5
    alias: node-5
  - id: node-6
    alias: node-6
  - id: node-7
    alias: node-7
  - id: node-8
    alias: node-8
  - id: node-9
    alias: node-9
  - id: node-10
    alias: node-10
`

func parseFixture(t *testing.T) map[string]interface{} {
	t.Helper()
	var doc map[string]interface{}
	if err := yaml.Unmarshal([]byte(fixtureYAML), &doc); err != nil {
		t.Fatalf("failed to parse fixture YAML: %v", err)
	}
	return doc
}

func TestQueryEnvironmentProjectionIsBounded(t *testing.T) {
	data := parseFixture(t)
	proj := buildEnvProjection(data)

	if proj.InventoryID != "inventory-dev" {
		t.Errorf("expected inventory_id=inventory-dev, got %q", proj.InventoryID)
	}
	if proj.CommonInventory != "agent-test/inventory.yaml" {
		t.Errorf("expected common_inventory=agent-test/inventory.yaml, got %q", proj.CommonInventory)
	}
	if proj.Env != "dev" {
		t.Errorf("expected env=dev, got %q", proj.Env)
	}
	if proj.Profile != "dev-runtime-provider-pool" {
		t.Errorf("expected profile=dev-runtime-provider-pool, got %q", proj.Profile)
	}
	if proj.LastUpdatedAt != "2026-07-13" {
		t.Errorf("expected last_updated_at=2026-07-13, got %q", proj.LastUpdatedAt)
	}
	if proj.Source == nil {
		t.Error("expected source to be present in projection")
	}
	if proj.Edge == nil {
		t.Error("expected edge to be present in projection")
	}
	if proj.Build == nil {
		t.Error("expected build to be present in projection")
	}

	// Ensure model and nodes are NOT in the projection.
	jsonBytes, err := json.Marshal(proj)
	if err != nil {
		t.Fatalf("failed to marshal projection: %v", err)
	}
	jsonStr := string(jsonBytes)
	if strings.Contains(jsonStr, `"model"`) {
		t.Error("projection should not contain model key")
	}
	if strings.Contains(jsonStr, `"nodes"`) {
		t.Error("projection should not contain nodes key")
	}
}

func TestQueryNodeByIDAndAlias(t *testing.T) {
	data := parseFixture(t)

	// Match by id.
	matches := queryNodes(data, "gx10-vllm-node")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match for id=gx10-vllm-node, got %d: %+v", len(matches), matches)
	}
	if matches[0].Path != "nodes[1]" {
		t.Errorf("expected path=nodes[1], got %q", matches[0].Path)
	}
	nodeObj, ok := matches[0].Value.(map[string]interface{})
	if !ok {
		t.Fatalf("expected map[string]interface{}, got %T", matches[0].Value)
	}
	if nodeObj["id"] != "gx10-vllm-node" {
		t.Errorf("expected id=gx10-vllm-node, got %v", nodeObj["id"])
	}

	// Match by alias.
	matches = queryNodes(data, "mac-codex")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match for alias=mac-codex, got %d: %+v", len(matches), matches)
	}
	if matches[0].Path != "nodes[0]" {
		t.Errorf("expected path=nodes[0], got %q", matches[0].Path)
	}
}

func TestQueryModelAcrossAliasShapes(t *testing.T) {
	data := parseFixture(t)

	// Match by alias key in model.aliases.
	matches := queryModels(data, "gemma4:26b")
	if len(matches) == 0 {
		t.Fatal("expected at least 1 match for alias key gemma4:26b")
	}
	found := false
	for _, m := range matches {
		if m.Path == "model.aliases.gemma4:26b" {
			found = true
			modelMap, ok := m.Value.(map[string]interface{})
			if !ok {
				t.Fatalf("expected map[string]interface{}, got %T", m.Value)
			}
			if modelMap["capacity_total"] != 5 {
				t.Errorf("expected capacity_total=5, got %v", modelMap["capacity_total"])
			}
			break
		}
	}
	if !found {
		t.Errorf("expected match at model.aliases.gemma4:26b, got %+v", matches)
	}

	// Match by entity alias in model subtree.
	matches = queryModels(data, "qwen3.6:35b")
	if len(matches) == 0 {
		t.Fatal("expected at least 1 match for entity alias qwen3.6:35b")
	}
	found = false
	for _, m := range matches {
		if m.Path == "model.qwen3_6_reference.alias" && m.Value == "qwen3.6:35b" {
			found = true
			break
		}
	}
	if !found {
		t.Errorf("expected match at model.qwen3_6_reference.alias, got %+v", matches)
	}

	// Match by entity id in model subtree.
	matches = queryModels(data, "qwen3-6-ref")
	if len(matches) == 0 {
		t.Fatal("expected at least 1 match for entity id qwen3-6-ref")
	}
	found = false
	for _, m := range matches {
		if m.Path == "model.qwen3_6_reference.id" && m.Value == "qwen3-6-ref" {
			found = true
			break
		}
	}
	if !found {
		t.Errorf("expected match at model.qwen3_6_reference.id, got %+v", matches)
	}
}

func TestQueryProviderSortsAndDeduplicatesMatches(t *testing.T) {
	data := parseFixture(t)

	matches := queryProviders(data, "mac-mlx-vllm")
	if len(matches) == 0 {
		t.Fatal("expected matches for provider mac-mlx-vllm")
	}

	// Check no duplicates.
	paths := make(map[string]bool)
	for _, m := range matches {
		if paths[m.Path] {
			t.Errorf("duplicate path found: %s", m.Path)
		}
		paths[m.Path] = true
	}

	// Check sorted by path.
	for i := 1; i < len(matches); i++ {
		if matches[i].Path < matches[i-1].Path {
			t.Errorf("matches not sorted: %s < %s", matches[i].Path, matches[i-1].Path)
		}
	}

	// Verify expected paths exist.
	expectedPaths := map[string]bool{
		"nodes[0].providers[0]":                 false,
		"model.aliases.gemma4:26b.providers[0]": false,
		"model.qwen3_6_reference.provider":      false,
	}
	for _, m := range matches {
		if _, ok := expectedPaths[m.Path]; ok {
			expectedPaths[m.Path] = true
		}
	}
	for p, found := range expectedPaths {
		if !found {
			t.Errorf("expected path %s not found in matches %+v", p, matches)
		}
	}
}

func TestQueryZeroMatch(t *testing.T) {
	data := parseFixture(t)

	matches := queryNodes(data, "nonexistent")
	if len(matches) != 0 {
		t.Errorf("expected 0 matches, got %d: %+v", len(matches), matches)
	}

	matches = queryProviders(data, "nonexistent-provider")
	if len(matches) != 0 {
		t.Errorf("expected 0 matches, got %d: %+v", len(matches), matches)
	}

	matches = queryModels(data, "nonexistent-model")
	if len(matches) != 0 {
		t.Errorf("expected 0 matches, got %d: %+v", len(matches), matches)
	}
}

func TestParseFlagsRejectsInvalidCombinations(t *testing.T) {
	tests := []struct {
		name    string
		args    []string
		wantErr bool
	}{
		{"missing env", []string{"--model", "x"}, true},
		{"multiple selectors", []string{"--env", "dev", "--model", "x", "--node", "y"}, true},
		{"unknown flag", []string{"--env", "dev", "--unknown"}, true},
		{"env without value", []string{"--env"}, true},
		{"model without value", []string{"--env", "dev", "--model"}, true},
		{"invalid env value", []string{"--env", "invalid-env"}, true},
		{"valid no selector", []string{"--env", "dev"}, false},
		{"valid single selector", []string{"--env", "dev", "--node", "x"}, false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := parseFlags(tt.args)
			if tt.wantErr && err == nil {
				t.Error("expected error but got nil")
			}
			if !tt.wantErr && err != nil {
				t.Errorf("unexpected error: %v", err)
			}
		})
	}
}

func TestQueryIsDeterministic(t *testing.T) {
	data := parseFixture(t)

	var outputs []string
	for i := 0; i < 5; i++ {
		matches := queryProviders(data, "gx10-vllm")
		out, err := json.Marshal(matches)
		if err != nil {
			t.Fatalf("json marshal error: %v", err)
		}
		outputs = append(outputs, string(out))
	}

	for i := 1; i < len(outputs); i++ {
		if outputs[i] != outputs[0] {
			t.Errorf("query not deterministic: output %d differs from output 0", i)
		}
	}
}

func TestQueryProviderInDirectProviders(t *testing.T) {
	data := parseFixture(t)

	matches := queryProviders(data, "ornith-direct")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match for ornith-direct, got %d: %+v", len(matches), matches)
	}
	if matches[0].Path != "nodes[2].direct_providers[0]" {
		t.Errorf("expected path=nodes[2].direct_providers[0], got %q", matches[0].Path)
	}
}

func TestQueryModelAliasKey(t *testing.T) {
	data := parseFixture(t)

	matches := queryModels(data, "ornith:35b")
	if len(matches) == 0 {
		t.Fatal("expected matches for ornith:35b")
	}

	foundAliasKey := false
	foundEntity := false
	for _, m := range matches {
		if m.Path == "model.aliases.ornith:35b" {
			foundAliasKey = true
		}
		if m.Path == "model.alias" && m.Value == "ornith:35b" {
			foundEntity = true
		}
	}
	if !foundAliasKey {
		t.Errorf("expected match at model.aliases.ornith:35b, got %+v", matches)
	}
	if !foundEntity {
		t.Errorf("expected match at model.alias, got %+v", matches)
	}
}

func TestQueryProviderInModelAliasesProviders(t *testing.T) {
	data := parseFixture(t)

	matches := queryProviders(data, "corp-dgx-spark-01-ornith")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match for corp-dgx-spark-01-ornith, got %d: %+v", len(matches), matches)
	}
	if matches[0].Path != "model.aliases.ornith:35b.providers[0]" {
		t.Errorf("expected path=model.aliases.ornith:35b.providers[0], got %q", matches[0].Path)
	}
}

func TestEncodeJSONOutput(t *testing.T) {
	var buf bytes.Buffer
	proj := buildEnvProjection(parseFixture(t))
	encodeJSON(&buf, proj)

	if buf.Len() == 0 {
		t.Error("expected non-empty JSON output")
	}

	var result envProjection
	if err := json.Unmarshal(buf.Bytes(), &result); err != nil {
		t.Fatalf("failed to unmarshal JSON output: %v\noutput: %s", err, buf.String())
	}
	if result.Env != "dev" {
		t.Errorf("expected env=dev in JSON output, got %q", result.Env)
	}
}

func TestQueryNodeReturnsEntityValue(t *testing.T) {
	data := parseFixture(t)
	matches := queryNodes(data, "gx10-vllm-node")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match, got %d", len(matches))
	}
	nodeMap, ok := matches[0].Value.(map[string]interface{})
	if !ok {
		t.Fatalf("expected map[string]interface{}, got %T", matches[0].Value)
	}
	if nodeMap["id"] != "gx10-vllm-node" {
		t.Errorf("expected id=gx10-vllm-node, got %v", nodeMap["id"])
	}
}

func TestQueryModelAliasReturnsValueAtPath(t *testing.T) {
	data := parseFixture(t)
	matches := queryModels(data, "gemma4:26b")
	if len(matches) != 1 {
		t.Fatalf("expected 1 match, got %d", len(matches))
	}
	modelMap, ok := matches[0].Value.(map[string]interface{})
	if !ok {
		t.Fatalf("expected map[string]interface{}, got %T", matches[0].Value)
	}
	if modelMap["capacity_total"] != 5 {
		t.Errorf("expected capacity_total=5, got %v", modelMap["capacity_total"])
	}
}

func TestQueryProviderAcrossCanonicalShapes(t *testing.T) {
	data := parseFixture(t)
	matches := queryProviders(data, "gx10-vllm")
	if len(matches) != 2 {
		t.Fatalf("expected 2 matches, got %d: %+v", len(matches), matches)
	}

	paths := map[string]interface{}{
		matches[0].Path: matches[0].Value,
		matches[1].Path: matches[1].Value,
	}

	val1, ok := paths["model.active_edge_model_group.providers.gx10-vllm"]
	if !ok {
		t.Errorf("expected path model.active_edge_model_group.providers.gx10-vllm, got keys: %v", paths)
	} else {
		m, ok := val1.(map[string]interface{})
		if !ok || m["capacity_total"] != 4 {
			t.Errorf("expected map with capacity_total=4, got %T: %v", val1, val1)
		}
	}

	val2, ok := paths["nodes[1].provider"]
	if !ok {
		t.Errorf("expected path nodes[1].provider, got keys: %v", paths)
	} else {
		m, ok := val2.(map[string]interface{})
		if !ok || m["id"] != "gx10-vllm" {
			t.Errorf("expected map with id=gx10-vllm, got %T: %v", val2, val2)
		}
	}
}

func TestQueryNodeSortsByPath(t *testing.T) {
	data := parseFixture(t)
	nodes, ok := data["nodes"].([]interface{})
	if !ok {
		t.Fatalf("nodes slice not found")
	}
	for i, n := range nodes {
		obj, ok := n.(map[string]interface{})
		if ok {
			obj["alias"] = "target-node"
			obj["id"] = fmt.Sprintf("node-id-%d", i)
		}
	}
	matches := queryNodes(data, "target-node")
	if len(matches) < 10 {
		t.Fatalf("expected at least 10 matches, got %d", len(matches))
	}
	for i := 1; i < len(matches); i++ {
		if matches[i].Path < matches[i-1].Path {
			t.Errorf("matches not sorted by path: %s < %s", matches[i].Path, matches[i-1].Path)
		}
	}
}

func TestParseFlagsRejectsUnsupportedEnv(t *testing.T) {
	_, err := parseFlags([]string{"--env", "unsupported"})
	if err == nil {
		t.Error("expected error for unsupported env, got nil")
	}
}

func TestValidateInventoryRejectsIdentityMismatch(t *testing.T) {
	data := parseFixture(t)
	err := validateInventory(data, "dev-corp", "")
	if err == nil {
		t.Error("expected identity mismatch error, got nil")
	}
}

func TestResolveInventoryPathFromSharedManifest(t *testing.T) {
	manifest := []byte(`
inventory_id: inventory
environments:
  dev:
    inventory_id: inventory-dev
    path: agent-test/inventory-dev.yaml
  dev-corp:
    inventory_id: inventory-dev-corp
    path: agent-test/inventory-dev-corp.yaml
`)

	path, err := resolveInventoryPath(manifest, "dev-corp")
	if err != nil {
		t.Fatalf("unexpected resolve error: %v", err)
	}
	if path != "agent-test/inventory-dev-corp.yaml" {
		t.Fatalf("expected dev-corp inventory path, got %q", path)
	}
}

func TestResolveInventoryPathRejectsIdentityMismatch(t *testing.T) {
	manifest := []byte(`
inventory_id: inventory
environments:
  dev:
    inventory_id: inventory-dev-corp
    path: agent-test/inventory-dev.yaml
`)

	if _, err := resolveInventoryPath(manifest, "dev"); err == nil {
		t.Fatal("expected inventory identity mismatch")
	}
}

func TestValidateInventoryRejectsInvalidProjectionTypes(t *testing.T) {
	data := parseFixture(t)
	delete(data, "profile")
	err := validateInventory(data, "dev", "")
	if err == nil {
		t.Error("expected error for missing profile, got nil")
	}
}

func TestValidateInventoryRequiresCommonInventoryIdentity(t *testing.T) {
	data := parseFixture(t)
	delete(data, "common_inventory")
	if err := validateInventory(data, "dev", ""); err == nil {
		t.Fatal("expected common_inventory validation error")
	}
}

func TestRunClassifiesSchemaError(t *testing.T) {
	data := parseFixture(t)
	delete(data, "profile")

	var stdout, stderr bytes.Buffer
	loader := func(env string) (map[string]interface{}, error) {
		return data, nil
	}

	code := run([]string{"--env", "dev"}, &stdout, &stderr, loader)
	if code != exitError {
		t.Errorf("expected exitError(2), got %d. stderr: %s", code, stderr.String())
	}
}

func TestRunExitCodesAndStreams(t *testing.T) {
	tests := []struct {
		name       string
		args       []string
		modifyData func(map[string]interface{})
		wantCode   int
		wantStdout bool
		wantStderr bool
	}{
		{
			name:       "success env projection",
			args:       []string{"--env", "dev"},
			wantCode:   exitOK,
			wantStdout: true,
			wantStderr: false,
		},
		{
			name:       "success query match",
			args:       []string{"--env", "dev", "--node", "gx10-vllm-node"},
			wantCode:   exitOK,
			wantStdout: true,
			wantStderr: false,
		},
		{
			name:       "zero match",
			args:       []string{"--env", "dev", "--node", "definitely-not-a-node"},
			wantCode:   exitZeroMatch,
			wantStdout: false,
			wantStderr: true,
		},
		{
			name:       "flag error",
			args:       []string{"--env", "dev", "--model", "x", "--node", "y"},
			wantCode:   exitError,
			wantStdout: false,
			wantStderr: true,
		},
		{
			name:       "schema error (missing edge)",
			args:       []string{"--env", "dev"},
			modifyData: func(d map[string]interface{}) { delete(d, "edge") },
			wantCode:   exitError,
			wantStdout: false,
			wantStderr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			data := parseFixture(t)
			if tt.modifyData != nil {
				tt.modifyData(data)
			}
			loader := func(env string) (map[string]interface{}, error) {
				return data, nil
			}
			var stdout, stderr bytes.Buffer
			code := run(tt.args, &stdout, &stderr, loader)
			if code != tt.wantCode {
				t.Errorf("expected code %d, got %d. stderr: %s", tt.wantCode, code, stderr.String())
			}
			if tt.wantStdout && stdout.Len() == 0 {
				t.Error("expected output in stdout, got empty")
			}
			if !tt.wantStdout && stdout.Len() > 0 {
				t.Errorf("expected empty stdout, got %q", stdout.String())
			}
			if tt.wantStderr && stderr.Len() == 0 {
				t.Error("expected output in stderr, got empty")
			}
			if !tt.wantStderr && stderr.Len() > 0 {
				t.Errorf("expected empty stderr, got %q", stderr.String())
			}
		})
	}
}

func TestValidateInventoryRejectsInvalidProjectionObjectTypes(t *testing.T) {
	tests := []struct {
		name    string
		modify  func(map[string]interface{})
		wantErr string
	}{
		{
			name:    "source is scalar",
			modify:  func(d map[string]interface{}) { d["source"] = "invalid" },
			wantErr: "missing or invalid type for source",
		},
		{
			name:    "source is list",
			modify:  func(d map[string]interface{}) { d["source"] = []interface{}{"a"} },
			wantErr: "missing or invalid type for source",
		},
		{
			name:    "edge is null",
			modify:  func(d map[string]interface{}) { d["edge"] = nil },
			wantErr: "missing or invalid type for edge",
		},
		{
			name:    "build is scalar",
			modify:  func(d map[string]interface{}) { d["build"] = 42 },
			wantErr: "missing or invalid type for build",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			data := parseFixture(t)
			tt.modify(data)
			err := validateInventory(data, "dev", "")
			if err == nil {
				t.Error("expected error, got nil")
			} else if err.Error() != tt.wantErr {
				t.Errorf("expected error %q, got %q", tt.wantErr, err.Error())
			}
		})
	}
}

func TestValidateInventoryRequiresProviderModelAndNodes(t *testing.T) {
	data := parseFixture(t)

	// Remove model for provider query.
	delete(data, "model")
	err := validateInventory(data, "dev", "provider")
	if err == nil {
		t.Error("expected error for missing model in provider query, got nil")
	} else if err.Error() != "missing or invalid type for model" {
		t.Errorf("expected model error, got %q", err.Error())
	}

	// Restore model but remove nodes for provider query.
	data["model"] = map[string]interface{}{
		"aliases": map[string]interface{}{},
	}
	delete(data, "nodes")
	err = validateInventory(data, "dev", "provider")
	if err == nil {
		t.Error("expected error for missing nodes in provider query, got nil")
	} else if err.Error() != "missing or invalid type for nodes" {
		t.Errorf("expected nodes error, got %q", err.Error())
	}

	// nodes is wrong type for provider query.
	data["nodes"] = "invalid"
	err = validateInventory(data, "dev", "provider")
	if err == nil {
		t.Error("expected error for invalid nodes type in provider query, got nil")
	}

	// nodes is wrong type for node query.
	delete(data, "model")
	err = validateInventory(data, "dev", "node")
	if err == nil {
		t.Error("expected error for invalid nodes type in node query, got nil")
	}
}

func TestRunClassifiesProviderSchemaError(t *testing.T) {
	data := parseFixture(t)
	delete(data, "model")

	var stdout, stderr bytes.Buffer
	loader := func(env string) (map[string]interface{}, error) {
		return data, nil
	}

	code := run([]string{"--env", "dev", "--provider", "test"}, &stdout, &stderr, loader)
	if code != exitError {
		t.Errorf("expected exitError(2), got %d. stderr: %s", code, stderr.String())
	}
	if stdout.Len() > 0 {
		t.Errorf("expected empty stdout, got %q", stdout.String())
	}
	if stderr.Len() == 0 {
		t.Error("expected error message in stderr")
	}
}

func TestRunHelpUsesInjectedStreams(t *testing.T) {
	tests := []struct {
		name string
		args []string
	}{
		{"help flag", []string{"--env", "dev", "--help"}},
		{"help short", []string{"--env", "dev", "-h"}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			data := parseFixture(t)
			var stdout, stderr bytes.Buffer
			loaderCalled := false
			loader := func(env string) (map[string]interface{}, error) {
				loaderCalled = true
				return data, nil
			}

			code := run(tt.args, &stdout, &stderr, loader)
			if code != exitOK {
				t.Errorf("expected exitOK(0), got %d", code)
			}
			if loaderCalled {
				t.Error("loader should not be called for help")
			}
			if stdout.Len() == 0 {
				t.Error("expected usage output in stdout")
			}
			if !strings.Contains(stdout.String(), "usage:") {
				t.Errorf("expected usage text in stdout, got %q", stdout.String())
			}
			if stderr.Len() > 0 {
				t.Errorf("expected empty stderr, got %q", stderr.String())
			}
		})
	}
}