Function TestGPUCounts

utils/spec_store_test.go:47–78

(t *testing.T)

}

func TestGPUCounts(t *testing.T) {
	s := testSpecStore()

	tests := []struct {
		name     string
		gpuType  string
		expected []int
	}{
		{
			name:     "a6000 only has 1 GPU",
			gpuType:  "a6000",
			expected: []int{1},
		},
		{
			name:     "a100xl has 1, 2, 4, and 8 GPUs",
			gpuType:  "a100xl",
			expected: []int{1, 2, 4, 8},
		},
		{
			name:     "unknown GPU returns nil",
			gpuType:  "unknown",
			expected: nil,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := s.GPUCounts(tt.gpuType)
			assert.Equal(t, tt.expected, got)
		})
	}
}

func TestNeedsGPUCountPhase(t *testing.T) {
	s := testSpecStore()