forge/pkg/imager/pull_test.go

package imager

import (
	"context"
	"os"
	"path/filepath"
	"testing"
	"time"

	"yizhisec.com/hsv2/forge/pkg/logger"
)

// TestPullImage_PublicImage tests pulling a public image from Docker Hub
func TestPullImage_PublicImage(t *testing.T) {
	logger.SetLogLevel(logger.LogLevelDebug)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	// Create temp directory for test output
	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "alpine.tar")
	outputFile = "./redis.alpine.tar"

	// Create output file
	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Pull alpine image (small and commonly available)
	imageName := "docker-mirror.yizhisec.com/library/redis:alpine"
	err = PullImage(ctx, imageName, f)
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	// Verify file was created and has content
	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s, size: %d bytes", imageName, info.Size())
}

// TestPullImage_WithRegistry tests pulling from a specific registry
func TestPullImage_WithRegistry(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "nginx.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Pull from docker.io with explicit registry
	imageName := "docker.io/library/nginx:alpine"
	err = PullImage(ctx, imageName, f)
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s, size: %d bytes", imageName, info.Size())
}

// TestPullImage_WithRename tests pulling and renaming an image
func TestPullImage_WithRename(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "renamed.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	imageName := "alpine:3.19"
	newName := "my-custom-alpine:latest"

	err = PullImage(ctx, imageName, f, WithRename(newName))
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s and renamed to %s, size: %d bytes", imageName, newName, info.Size())
}

// TestPullImage_WithAuth tests pulling with authentication (skip if no credentials)
func TestPullImage_WithAuth(t *testing.T) {
	// This test requires actual credentials, so it's skipped by default
	t.Skip("Skipping authentication test - requires valid credentials")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "private.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Replace with your private registry details
	imageName := "your-registry.com/your-repo/your-image:tag"
	username := "your-username"
	password := "your-password"

	err = PullImage(ctx, imageName, f, WithAuth(username, password))
	if err != nil {
		t.Fatalf("failed to pull private image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled private image %s, size: %d bytes", imageName, info.Size())
}

// TestPullImage_WithRetry tests retry functionality
func TestPullImage_WithRetry(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "retry.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Pull with custom retry settings
	imageName := "alpine:3.19"
	err = PullImage(ctx, imageName, f,
		WithRetry(5),                  // 5 retries
		WithRetryDelay(1*time.Second), // 1 second delay
	)
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s with retry, size: %d bytes", imageName, info.Size())
}

// TestPullImage_WithZeroRetry tests with retry disabled
func TestPullImage_WithZeroRetry(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "no-retry.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Pull with retry disabled
	imageName := "alpine:3.19"
	err = PullImage(ctx, imageName, f, WithRetry(0))
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s without retry, size: %d bytes", imageName, info.Size())
}

// TestPullImage_WithSkipTLSVerify tests pulling with TLS verification disabled
func TestPullImage_WithSkipTLSVerify(t *testing.T) {
	// This test is for registries with self-signed certificates
	t.Skip("Skipping TLS skip test - requires a registry with self-signed cert")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "insecure.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	imageName := "your-insecure-registry.local/image:tag"
	err = PullImage(ctx, imageName, f, WithSkipTLSVerify())
	if err != nil {
		t.Fatalf("failed to pull image %s: %v", imageName, err)
	}

	if err := f.Close(); err != nil {
		t.Fatalf("failed to close file: %v", err)
	}

	info, err := os.Stat(outputFile)
	if err != nil {
		t.Fatalf("failed to stat output file: %v", err)
	}

	if info.Size() == 0 {
		t.Error("output file is empty")
	}

	t.Logf("Successfully pulled %s with TLS verification skipped, size: %d bytes", imageName, info.Size())
}

// TestParseImageReference tests the image reference parsing logic
func TestParseImageReference(t *testing.T) {
	tests := []struct {
		name           string
		imageName      string
		wantRegistry   string
		wantRepository string
		wantTag        string
		wantErr        bool
	}{
		{
			name:           "simple image with tag",
			imageName:      "alpine:3.19",
			wantRegistry:   "registry-1.docker.io",
			wantRepository: "library/alpine",
			wantTag:        "3.19",
			wantErr:        false,
		},
		{
			name:           "simple image without tag",
			imageName:      "alpine",
			wantRegistry:   "registry-1.docker.io",
			wantRepository: "library/alpine",
			wantTag:        "latest",
			wantErr:        false,
		},
		{
			name:           "image with namespace",
			imageName:      "library/nginx:alpine",
			wantRegistry:   "registry-1.docker.io",
			wantRepository: "library/nginx",
			wantTag:        "alpine",
			wantErr:        false,
		},
		{
			name:           "custom registry",
			imageName:      "hub.yizhisec.com/external/alpine:3.22.2",
			wantRegistry:   "hub.yizhisec.com",
			wantRepository: "external/alpine",
			wantTag:        "3.22.2",
			wantErr:        false,
		},
		{
			name:           "docker.io explicit",
			imageName:      "docker.io/library/redis:8.2.2",
			wantRegistry:   "docker.io",
			wantRepository: "library/redis",
			wantTag:        "8.2.2",
			wantErr:        false,
		},
		{
			name:           "registry with port",
			imageName:      "localhost:5000/myimage:v1.0",
			wantRegistry:   "localhost:5000",
			wantRepository: "myimage",
			wantTag:        "v1.0",
			wantErr:        false,
		},
		{
			name:           "quay.io registry",
			imageName:      "quay.io/k0sproject/pause:3.10.1",
			wantRegistry:   "quay.io",
			wantRepository: "k0sproject/pause",
			wantTag:        "3.10.1",
			wantErr:        false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ref, err := parseImageReference(tt.imageName)
			if (err != nil) != tt.wantErr {
				t.Errorf("parseImageReference() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if err != nil {
				return
			}

			if ref.Registry != tt.wantRegistry {
				t.Errorf("parseImageReference() registry = %v, want %v", ref.Registry, tt.wantRegistry)
			}
			if ref.Repository != tt.wantRepository {
				t.Errorf("parseImageReference() repository = %v, want %v", ref.Repository, tt.wantRepository)
			}
			if ref.Tag != tt.wantTag {
				t.Errorf("parseImageReference() tag = %v, want %v", ref.Tag, tt.wantTag)
			}
		})
	}
}

// TestPullImage_ContextCancellation tests that context cancellation works
func TestPullImage_ContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	tmpDir := t.TempDir()
	outputFile := filepath.Join(tmpDir, "cancelled.tar")

	f, err := os.Create(outputFile)
	if err != nil {
		t.Fatalf("failed to create output file: %v", err)
	}
	defer f.Close()

	// Cancel context immediately
	cancel()

	imageName := "alpine:3.19"
	err = PullImage(ctx, imageName, f)
	if err == nil {
		t.Error("expected error due to cancelled context, got nil")
	}

	t.Logf("Context cancellation correctly handled: %v", err)
}

// BenchmarkPullImage benchmarks image pulling performance
func BenchmarkPullImage(b *testing.B) {
	ctx := context.Background()
	tmpDir := b.TempDir()

	imageName := "alpine:3.19"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		outputFile := filepath.Join(tmpDir, "alpine_bench.tar")
		f, err := os.Create(outputFile)
		if err != nil {
			b.Fatalf("failed to create output file: %v", err)
		}

		err = PullImage(ctx, imageName, f)
		if err != nil {
			b.Fatalf("failed to pull image: %v", err)
		}

		f.Close()
		os.Remove(outputFile)
	}
}