tyto/backend/internal/collectors/gpu/manager_test.go

package gpu

import (
	"os"
	"path/filepath"
	"testing"
)

func TestManager_NoGPUs(t *testing.T) {
	tmpDir := t.TempDir()
	manager := NewManager(tmpDir)

	if manager.Available() {
		t.Error("Expected no GPUs available")
	}

	if manager.GPUCount() != 0 {
		t.Errorf("Expected 0 GPUs, got %d", manager.GPUCount())
	}

	stats, err := manager.Collect()
	if err != nil {
		t.Fatalf("Collect failed: %v", err)
	}

	if stats.Available {
		t.Error("Expected stats.Available to be false")
	}

	if len(stats.GPUs) != 0 {
		t.Errorf("Expected 0 GPUs in stats, got %d", len(stats.GPUs))
	}
}

func TestManager_WithAMDGPU(t *testing.T) {
	tmpDir := t.TempDir()
	sysPath := filepath.Join(tmpDir, "sys")

	// Create mock AMD GPU
	gpuPath := filepath.Join(sysPath, "class/drm/card0/device")
	if err := os.MkdirAll(gpuPath, 0755); err != nil {
		t.Fatal(err)
	}

	driverTarget := filepath.Join(tmpDir, "drivers/amdgpu")
	if err := os.MkdirAll(driverTarget, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(driverTarget, filepath.Join(gpuPath, "driver")); err != nil {
		t.Fatal(err)
	}

	if err := os.WriteFile(filepath.Join(gpuPath, "gpu_busy_percent"), []byte("42\n"), 0644); err != nil {
		t.Fatal(err)
	}

	manager := NewManager(sysPath)

	if !manager.Available() {
		t.Error("Expected GPU to be available")
	}

	if manager.GPUCount() != 1 {
		t.Errorf("Expected 1 GPU, got %d", manager.GPUCount())
	}

	stats, err := manager.Collect()
	if err != nil {
		t.Fatalf("Collect failed: %v", err)
	}

	if !stats.Available {
		t.Error("Expected stats.Available to be true")
	}

	if len(stats.GPUs) != 1 {
		t.Fatalf("Expected 1 GPU in stats, got %d", len(stats.GPUs))
	}

	if stats.GPUs[0].Index != 0 {
		t.Errorf("Expected index 0, got %d", stats.GPUs[0].Index)
	}

	if stats.GPUs[0].Utilization != 42 {
		t.Errorf("Expected utilization 42, got %d", stats.GPUs[0].Utilization)
	}
}

func TestManager_MixedVendors(t *testing.T) {
	tmpDir := t.TempDir()
	sysPath := filepath.Join(tmpDir, "sys")

	// Create AMD GPU (card0)
	amdPath := filepath.Join(sysPath, "class/drm/card0/device")
	if err := os.MkdirAll(amdPath, 0755); err != nil {
		t.Fatal(err)
	}
	amdDriver := filepath.Join(tmpDir, "drivers/amdgpu")
	if err := os.MkdirAll(amdDriver, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(amdDriver, filepath.Join(amdPath, "driver")); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(amdPath, "gpu_busy_percent"), []byte("50\n"), 0644); err != nil {
		t.Fatal(err)
	}

	// Create Intel GPU (card1)
	intelPath := filepath.Join(sysPath, "class/drm/card1/device")
	if err := os.MkdirAll(intelPath, 0755); err != nil {
		t.Fatal(err)
	}
	intelDriver := filepath.Join(tmpDir, "drivers/i915")
	if err := os.MkdirAll(intelDriver, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(intelDriver, filepath.Join(intelPath, "driver")); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(intelPath, "gt_cur_freq_mhz"), []byte("1000\n"), 0644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(intelPath, "gt_max_freq_mhz"), []byte("1500\n"), 0644); err != nil {
		t.Fatal(err)
	}

	manager := NewManager(sysPath)

	if !manager.Available() {
		t.Error("Expected GPUs to be available")
	}

	if manager.GPUCount() != 2 {
		t.Errorf("Expected 2 GPUs, got %d", manager.GPUCount())
	}

	stats, err := manager.Collect()
	if err != nil {
		t.Fatalf("Collect failed: %v", err)
	}

	if len(stats.GPUs) != 2 {
		t.Fatalf("Expected 2 GPUs in stats, got %d", len(stats.GPUs))
	}

	// Verify indices are sequential
	for i, gpu := range stats.GPUs {
		if gpu.Index != i {
			t.Errorf("GPU %d has index %d, expected %d", i, gpu.Index, i)
		}
	}
}