package agent

import (
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

func TestAShardPlacedOnAMIGInstanceIsLaunchedWithItsUUID(t *testing.T) {
	// The control plane sends device 1001: card 0's instance 1, by the number
	// placement gives it. Only this node knows the instance's UUID, which is
	// the one thing CUDA will open it by.
	a := &Agent{}
	a.inv.GPUs = []discovery.GPU{
		{Index: 1000, UUID: "MIG-aaa", MIG: &discovery.MIGDevice{ParentIndex: 0, Device: 0, Profile: "3g.40gb"}},
		{Index: 1001, UUID: "MIG-bbb", MIG: &discovery.MIGDevice{ParentIndex: 0, Device: 1, Profile: "3g.40gb"}},
		{Index: 1, UUID: "GPU-whole"},
	}

	spec := engines.ShardSpec{GPUIndices: []int{1001}}
	a.localizeSpec(&spec)
	if got := engines.NCCLEnv(spec)["CUDA_VISIBLE_DEVICES"]; got != "MIG-bbb" {
		t.Errorf("CUDA_VISIBLE_DEVICES=%q, want MIG-bbb", got)
	}

	// A whole card is given by index, as it always was.
	whole := engines.ShardSpec{GPUIndices: []int{1}}
	a.localizeSpec(&whole)
	if whole.GPUUUIDs != nil || engines.NCCLEnv(whole)["CUDA_VISIBLE_DEVICES"] != "1" {
		t.Errorf("a whole card was given as %v / %q", whole.GPUUUIDs, engines.NCCLEnv(whole)["CUDA_VISIBLE_DEVICES"])
	}

	// A device this node does not have is not guessed at.
	gone := engines.ShardSpec{GPUIndices: []int{1002}}
	a.localizeSpec(&gone)
	if gone.GPUUUIDs != nil {
		t.Errorf("an unknown device was named %v", gone.GPUUUIDs)
	}
}