package discovery

import (
	"encoding/json"
	"testing"
)

func TestJSONShape(t *testing.T) {
	inv := Inventory{
		Source: "nvidia-smi",
		GPUs:   []GPU{{Index: 0, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559, ComputeCapability: "9.0"}},
	}
	b, err := json.Marshal(inv)
	if err != nil {
		t.Fatal(err)
	}

	var out struct {
		Source string `json:"source"`
		GPUs   []struct {
			Index int    `json:"index"`
			Class string `json:"class"`
			Mem   int    `json:"memory_total_mb"`
		} `json:"gpus"`
	}
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	if out.Source != "nvidia-smi" || len(out.GPUs) != 1 {
		t.Fatalf("shape wrong: %s", b)
	}
	// The derived class is what decides which GPUs are interchangeable, so a
	// consumer must not have to recompute it.
	if out.GPUs[0].Class != "h100-80g" {
		t.Errorf("class=%q, should be included in the JSON", out.GPUs[0].Class)
	}
	if out.GPUs[0].Mem != 81559 {
		t.Errorf("memory_total_mb=%d", out.GPUs[0].Mem)
	}
}