fetch_ml/tests/unit/security/gpu_audit_test.go

package security

import (
	"bytes"
	"log/slog"
	"strings"
	"testing"

	"github.com/jfraeys/fetch_ml/internal/worker"
)

// TestGPUDetectionAudit verifies that GPU detection method is logged at startup
// for audit and reproducibility purposes.
func TestGPUDetectionAudit(t *testing.T) {
	tests := []struct {
		name       string
		gpuType    string
		wantMethod string
	}{
		{
			name:       "nvidia detection logs method",
			gpuType:    "nvidia",
			wantMethod: "config",
		},
		{
			name:       "apple detection logs method",
			gpuType:    "apple",
			wantMethod: "config",
		},
		{
			name:       "none detection logs method",
			gpuType:    "none",
			wantMethod: "config",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Capture log output
			var logBuf bytes.Buffer
			handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
				Level: slog.LevelInfo,
			})
			logger := slog.New(handler)

			// Create config with specified GPU type
			cfg := &worker.Config{
				GPUVendor: tt.gpuType,
			}

			// Perform GPU detection
			factory := &worker.GPUDetectorFactory{}
			result := factory.CreateDetectorWithInfo(cfg)

			// Log the detection info (this simulates startup logging)
			logger.Info("GPU detection completed",
				"gpu_type", result.Info.GPUType,
				"detection_method", result.Info.DetectionMethod,
				"configured_vendor", result.Info.ConfiguredVendor,
			)

			// Verify log output contains detection method
			logOutput := logBuf.String()
			if !strings.Contains(logOutput, "GPU detection completed") {
				t.Error("expected 'GPU detection completed' in log output")
			}
			if !strings.Contains(logOutput, string(result.Info.DetectionMethod)) {
				t.Errorf("expected detection method %q in log output", result.Info.DetectionMethod)
			}
			if !strings.Contains(logOutput, "detection_method=") {
				t.Error("expected 'detection_method=' field in log output")
			}
		})
	}

	t.Run("env override detection logged", func(t *testing.T) {
		// Set env var to trigger env-based detection
		t.Setenv("FETCH_ML_GPU_TYPE", "nvidia")

		var logBuf bytes.Buffer
		handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
			Level: slog.LevelInfo,
		})
		logger := slog.New(handler)

		// Create factory and detect
		factory := &worker.GPUDetectorFactory{}
		result := factory.CreateDetectorWithInfo(nil)

		// Log detection
		logger.Info("GPU detection completed",
			"gpu_type", result.Info.GPUType,
			"detection_method", result.Info.DetectionMethod,
			"env_override_type", result.Info.EnvOverrideType,
		)

		logOutput := logBuf.String()

		// Should log env override
		if !strings.Contains(logOutput, "env_override_type=nvidia") {
			t.Error("expected env override type in log output")
		}

		// Detection method should indicate env was used
		if result.Info.DetectionMethod != worker.DetectionSourceEnvType {
			t.Errorf("detection method = %v, want %v", result.Info.DetectionMethod, worker.DetectionSourceEnvType)
		}
	})

	t.Run("detection info fields populated", func(t *testing.T) {
		// Set both env vars
		t.Setenv("FETCH_ML_GPU_TYPE", "apple")
		t.Setenv("FETCH_ML_GPU_COUNT", "4")

		factory := &worker.GPUDetectorFactory{}
		result := factory.CreateDetectorWithInfo(nil)

		// Verify all expected fields are populated
		if result.Info.GPUType == "" {
			t.Error("GPUType field is empty")
		}
		if result.Info.ConfiguredVendor == "" {
			t.Error("ConfiguredVendor field is empty")
		}
		if result.Info.DetectionMethod == "" {
			t.Error("DetectionMethod field is empty")
		}
		if result.Info.EnvOverrideType != "apple" {
			t.Errorf("EnvOverrideType = %v, want 'apple'", result.Info.EnvOverrideType)
		}
		if result.Info.EnvOverrideCount != 4 {
			t.Errorf("EnvOverrideCount = %v, want 4", result.Info.EnvOverrideCount)
		}

		// Verify detection source is valid
		validSources := []worker.DetectionSource{
			worker.DetectionSourceEnvType,
			worker.DetectionSourceEnvCount,
			worker.DetectionSourceEnvBoth,
			worker.DetectionSourceConfig,
			worker.DetectionSourceAuto,
		}
		found := false
		for _, source := range validSources {
			if result.Info.DetectionMethod == source {
				found = true
				break
			}
		}
		if !found {
			t.Errorf("invalid detection method: %v", result.Info.DetectionMethod)
		}
	})
}

// TestGPUDetectionStructuredLogging verifies structured logging format for audit trails
func TestGPUDetectionStructuredLogging(t *testing.T) {
	var logBuf bytes.Buffer
	handler := slog.NewJSONHandler(&logBuf, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	})
	logger := slog.New(handler)

	// Simulate startup detection
	cfg := &worker.Config{GPUVendor: "nvidia"}
	factory := &worker.GPUDetectorFactory{}
	result := factory.CreateDetectorWithInfo(cfg)

	// Log with structured format for audit
	logger.Info("gpu_detection_startup",
		"event_type", "gpu_detection",
		"gpu_type", result.Info.GPUType,
		"detection_method", result.Info.DetectionMethod,
		"configured_vendor", result.Info.ConfiguredVendor,
	)

	logOutput := logBuf.String()

	// Verify JSON structure contains required fields
	if !strings.Contains(logOutput, `"event_type":"gpu_detection"`) {
		t.Error("expected event_type field in JSON log")
	}
	if !strings.Contains(logOutput, `"detection_method"`) {
		t.Error("expected detection_method field in JSON log")
	}
	if !strings.Contains(logOutput, `"configured_vendor"`) {
		t.Error("expected configured_vendor field in JSON log")
	}
}

// TestGPUDetectionAuditDisabled verifies behavior when detection fails
func TestGPUDetectionAuditDisabled(t *testing.T) {
	var logBuf bytes.Buffer
	handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	})
	logger := slog.New(handler)

	// Test with GPU none - should still log
	cfg := &worker.Config{GPUVendor: "none"}
	factory := &worker.GPUDetectorFactory{}
	result := factory.CreateDetectorWithInfo(cfg)

	logger.Info("GPU detection completed",
		"gpu_type", result.Info.GPUType,
		"detection_method", result.Info.DetectionMethod,
	)

	logOutput := logBuf.String()

	// Should still log even when GPU is none
	if !strings.Contains(logOutput, "GPU detection completed") {
		t.Error("expected log output even for GPU type 'none'")
	}

	// Should indicate config-based detection
	if result.Info.DetectionMethod != worker.DetectionSourceConfig {
		t.Errorf("detection method = %v, want %v", result.Info.DetectionMethod, worker.DetectionSourceConfig)
	}
}