AetherForge/agent/miner/container_launcher_test.go

package miner

import (
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"crypto-miner-agent/config"
)

func longRunningTestCmd() *exec.Cmd {
	if runtime.GOOS == "windows" {
		return exec.Command("ping", "-n", "600", "127.0.0.1")
	}
	return exec.Command("sleep", "600")
}

func quickExitTestCmd() *exec.Cmd {
	if runtime.GOOS == "windows" {
		return exec.Command("cmd", "/c", "exit", "0")
	}
	return exec.Command("true")
}

func dockerEnvFromArgs(args []string) map[string]string {
	out := make(map[string]string)
	for i := 0; i < len(args); i++ {
		if args[i] != "-e" || i+1 >= len(args) {
			continue
		}
		i++
		k, v, ok := strings.Cut(args[i], "=")
		if !ok {
			continue
		}
		out[k] = v
	}
	return out
}

func TestContainerLauncherStartWithFakeRuntime(t *testing.T) {
	const customImage = "registry.example/aether-worker:test"
	t.Setenv("AETHERFORGE_MINER_IMAGE", customImage)

	var gotCLI string
	var gotArgs []string
	SetContainerExecCommand(func(name string, args ...string) *exec.Cmd {
		if len(args) > 0 && args[0] == "rm" {
			return quickExitTestCmd()
		}
		gotCLI = name
		gotArgs = append([]string(nil), args...)
		return longRunningTestCmd()
	})
	defer SetContainerExecCommand(nil)

	cfg := config.RuntimeConfig{
		BuiltinConfig: config.BuiltinConfig{
			BuildID:    "test-build",
			ServerURL:  "http://c2.example",
			Wallet:     "XMR:wallet",
			WorkerName: "worker-1",
			PoolHost:   "pool.example.com",
			PoolPort:   3333,
			PoolTLS:    true,
			PoolPass:    "x",
			ThreadMode:  "fixed",
			Threads:     4,
			FleetSecret: "fleet-secret",
		},
	}
	rt := ContainerRuntimeInfo{Available: true, CLI: "docker", Version: "24.0.0"}

	launcher, err := NewContainerLauncher(cfg, rt)
	if err != nil {
		t.Fatalf("NewContainerLauncher: %v", err)
	}
	if launcher.Image() != customImage {
		t.Fatalf("Image()=%q want %q", launcher.Image(), customImage)
	}

	if err := launcher.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer launcher.Stop()

	if gotCLI != "docker" {
		t.Fatalf("runtime CLI=%q want docker", gotCLI)
	}

	wantPrefix := []string{"run", "--rm", "-d", "--name", "aetherforge-miner-test-build"}
	if len(gotArgs) < len(wantPrefix) {
		t.Fatalf("args=%v too short, want prefix %v", gotArgs, wantPrefix)
	}
	for i, w := range wantPrefix {
		if gotArgs[i] != w {
			t.Fatalf("args[%d]=%q want %q full=%v", i, gotArgs[i], w, gotArgs)
		}
	}

	if runtime.GOOS == "linux" {
		if !containsSeq(gotArgs, "--network", "host") {
			t.Fatalf("linux args missing --network host: %v", gotArgs)
		}
	}

	if gotArgs[len(gotArgs)-1] != customImage {
		t.Fatalf("image arg=%q want %q", gotArgs[len(gotArgs)-1], customImage)
	}

	env := dockerEnvFromArgs(gotArgs)
	wantEnv := map[string]string{
		"AETHERFORGE_SERVER_URL":      "http://c2.example",
		"AETHERFORGE_WALLET":            "XMR:wallet",
		"AETHERFORGE_WORKER":            "worker-1",
		"AETHERFORGE_POOL_HOST":         "pool.example.com",
		"AETHERFORGE_POOL_PORT":         "3333",
		"AETHERFORGE_POOL_TLS":          "1",
		"AETHERFORGE_POOL_PASS":         "x",
		"AETHERFORGE_THREADS":           "4",
		"AETHERFORGE_MINER_EXECUTION":   ExecutionInProcess,
		"AETHERFORGE_FLEET_SECRET":      "fleet-secret",
		"MINER_LOG_FILE":                  "/tmp/miner.log",
	}
	for k, want := range wantEnv {
		if got := env[k]; got != want {
			t.Fatalf("env[%s]=%q want %q", k, got, want)
		}
	}

	if !launcher.Running() {
		t.Fatal("Running() false after Start")
	}

	if err := launcher.Start(); err != nil {
		t.Fatalf("second Start: %v", err)
	}
	if !launcher.Running() {
		t.Fatal("Running() false after idempotent Start")
	}
}

func TestContainerLauncherDockerLoadFromTar(t *testing.T) {
	dir := t.TempDir()
	tarPath := filepath.Join(dir, "worker.tar")
	if err := os.WriteFile(tarPath, []byte("fake"), 0644); err != nil {
		t.Fatal(err)
	}

	var calls [][]string
	SetContainerExecCommand(func(name string, args ...string) *exec.Cmd {
		copied := append([]string{name}, args...)
		calls = append(calls, copied)
		if len(args) > 0 && args[0] == "load" {
			if runtime.GOOS == "windows" {
				return exec.Command("cmd", "/c", "echo Loaded image: aetherforge/agent-worker:tar")
			}
			return exec.Command("sh", "-c", "echo 'Loaded image: aetherforge/agent-worker:tar'")
		}
		if len(args) > 0 && args[0] == "rm" {
			return quickExitTestCmd()
		}
		return longRunningTestCmd()
	})
	defer SetContainerExecCommand(nil)

	cfg := config.RuntimeConfig{
		BuiltinConfig: config.BuiltinConfig{
			BuildID:  "tar-test",
			Wallet:   "XMR:wallet",
			PoolHost: "pool.example.com",
			PoolPort: 3333,
		},
	}
	rt := ContainerRuntimeInfo{Available: true, CLI: "docker"}

	launcher, err := NewContainerLauncherFromTar(cfg, rt, tarPath)
	if err != nil {
		t.Fatalf("NewContainerLauncherFromTar: %v", err)
	}
	if err := launcher.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer launcher.Stop()

	foundLoad := false
	foundNeverPull := false
	foundReadOnly := false
	for _, call := range calls {
		if len(call) >= 3 && call[1] == "load" && call[3] == tarPath {
			foundLoad = true
		}
		for i, arg := range call {
			if arg == "--pull=never" {
				foundNeverPull = true
			}
			if arg == "--read-only" {
				foundReadOnly = true
			}
			if arg == "--gpus" && i+1 < len(call) && call[i+1] == "all" {
				// gpu flag present when GPU enabled — optional in this cfg
			}
		}
	}
	if !foundLoad {
		t.Fatalf("docker load not invoked, calls=%v", calls)
	}
	if !foundNeverPull {
		t.Fatalf("expected --pull=never, calls=%v", calls)
	}
	if !foundReadOnly {
		t.Fatalf("expected --read-only, calls=%v", calls)
	}
}

func TestContainerLauncherPodmanStartWithFakeRuntime(t *testing.T) {
	var gotCLI string
	SetContainerExecCommand(func(name string, args ...string) *exec.Cmd {
		if len(args) > 0 && args[0] == "rm" {
			return quickExitTestCmd()
		}
		gotCLI = name
		return longRunningTestCmd()
	})
	defer SetContainerExecCommand(nil)

	cfg := config.RuntimeConfig{
		BuiltinConfig: config.BuiltinConfig{
			BuildID:  "podman-build",
			Wallet:   "XMR:wallet",
			PoolHost: "pool.example.com",
			PoolPort: 3333,
		},
	}
	rt := ContainerRuntimeInfo{Available: true, CLI: "podman", Version: "4.9.0"}

	launcher, err := NewContainerLauncher(cfg, rt)
	if err != nil {
		t.Fatalf("NewContainerLauncher: %v", err)
	}
	if err := launcher.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer launcher.Stop()

	if gotCLI != "podman" {
		t.Fatalf("runtime CLI=%q want podman", gotCLI)
	}
	if !launcher.Running() {
		t.Fatal("Running() false after podman Start")
	}
}

func TestContainerLauncherStopInvokesRm(t *testing.T) {
	var calls [][]string
	SetContainerExecCommand(func(name string, args ...string) *exec.Cmd {
		copied := append([]string{name}, args...)
		calls = append(calls, copied)
		if len(args) > 0 && args[0] == "rm" {
			return quickExitTestCmd()
		}
		return longRunningTestCmd()
	})
	defer SetContainerExecCommand(nil)

	cfg := config.RuntimeConfig{BuiltinConfig: config.BuiltinConfig{BuildID: "stop-test"}}
	rt := ContainerRuntimeInfo{Available: true, CLI: "docker"}
	launcher, err := NewContainerLauncher(cfg, rt)
	if err != nil {
		t.Fatal(err)
	}
	if err := launcher.Start(); err != nil {
		t.Fatal(err)
	}
	launcher.Stop()

	foundRm := false
	for _, call := range calls {
		if len(call) >= 3 && call[1] == "rm" && call[2] == "-f" {
			foundRm = true
			break
		}
	}
	if !foundRm {
		t.Fatalf("docker rm -f not invoked, calls=%v", calls)
	}
	if launcher.Running() {
		t.Fatal("Running() true after Stop")
	}
}

func TestContainerLauncherGPUFlags(t *testing.T) {
	var gotArgs []string
	SetContainerExecCommand(func(name string, args ...string) *exec.Cmd {
		if len(args) > 0 && args[0] == "rm" {
			return quickExitTestCmd()
		}
		gotArgs = append([]string(nil), args...)
		return longRunningTestCmd()
	})
	defer SetContainerExecCommand(nil)

	cfg := config.RuntimeConfig{
		BuiltinConfig: config.BuiltinConfig{
			BuildID:    "gpu-test",
			Wallet:     "XMR:wallet",
			PoolHost:   "pool.example.com",
			PoolPort:   3333,
			GPUEnabled: true,
		},
	}
	rt := ContainerRuntimeInfo{Available: true, CLI: "docker"}
	launcher, err := NewContainerLauncher(cfg, rt)
	if err != nil {
		t.Fatal(err)
	}
	if err := launcher.Start(); err != nil {
		t.Fatal(err)
	}
	defer launcher.Stop()

	if !containsSeq(gotArgs, "--gpus", "all") {
		t.Fatalf("args missing --gpus all: %v", gotArgs)
	}
}

func containsSeq(args []string, seq ...string) bool {
	if len(seq) == 0 || len(args) < len(seq) {
		return false
	}
	for i := 0; i <= len(args)-len(seq); i++ {
		match := true
		for j := range seq {
			if args[i+j] != seq[j] {
				match = false
				break
			}
		}
		if match {
			return true
		}
	}
	return false
}