metis/pkg/plan/workflow_test.go

package plan

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"testing"

	"metis/pkg/inventory"
)

func TestExecuteAndBuildImageFileWithFakes(t *testing.T) {
	rootTools := fakeRootfsTools(t)
	mountTools := fakeMountTools(t)
	t.Setenv("PATH", rootTools+string(os.PathListSeparator)+mountTools+string(os.PathListSeparator)+os.Getenv("PATH"))

	dir := t.TempDir()
	rawImage := filepath.Join(dir, "base.img")
	if err := os.WriteFile(rawImage, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	sum := imageChecksum(t, rawImage)
	inv := &inventory.Inventory{
		Classes: []inventory.NodeClass{{
			Name:     "c1",
			Arch:     "arm64",
			OS:       "linux",
			Image:    "file://" + rawImage,
			Checksum: sum,
		}},
		Nodes: []inventory.NodeSpec{{
			Name:          "n1",
			Class:         "c1",
			Hostname:      "n1",
			IP:            "10.0.0.1",
			K3sRole:       "agent",
			SSHUser:       "atlas",
			SSHAuthorized: []string{"ssh-ed25519 AAA"},
		}},
	}

	planDry, err := Execute(inv, "n1", filepath.Join(dir, "disk.img"), filepath.Join(dir, "cache"), false)
	if err != nil {
		t.Fatalf("Execute dry-run: %v", err)
	}
	if planDry.Node != "n1" || len(planDry.Actions) == 0 {
		t.Fatalf("unexpected dry-run plan: %#v", planDry)
	}

	bootDir := filepath.Join(dir, "boot")
	rootDir := filepath.Join(dir, "root")
	if err := os.MkdirAll(bootDir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(rootDir, 0o755); err != nil {
		t.Fatal(err)
	}
	t.Setenv("METIS_BOOT_PATH", bootDir)
	t.Setenv("METIS_ROOT_PATH", rootDir)
	t.Setenv("METIS_AUTO_MOUNT", "1")
	written := filepath.Join(dir, "written.img")
	planRun, err := Execute(inv, "n1", written, filepath.Join(dir, "cache2"), true)
	if err != nil {
		t.Fatalf("Execute confirm: %v", err)
	}
	if planRun.Image != "file://"+rawImage {
		t.Fatalf("unexpected plan image: %#v", planRun)
	}
	if _, err := os.Stat(written); err != nil {
		t.Fatalf("expected written image: %v", err)
	}
	if _, err := os.Stat(filepath.Join(rootDir, "etc/metis/firstboot.env")); err != nil {
		t.Fatalf("expected injected rootfs file: %v", err)
	}
}

func TestBuildImageFileMaterializesRootFS(t *testing.T) {
	rootTools := fakeRootfsTools(t)
	t.Setenv("PATH", rootTools+string(os.PathListSeparator)+os.Getenv("PATH"))

	dir := t.TempDir()
	rawImage := filepath.Join(dir, "base.img")
	if err := os.WriteFile(rawImage, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	sum := imageChecksum(t, rawImage)
	inv := &inventory.Inventory{
		Classes: []inventory.NodeClass{{
			Name:     "c1",
			Arch:     "arm64",
			OS:       "linux",
			Image:    "file://" + rawImage,
			Checksum: sum,
		}},
		Nodes: []inventory.NodeSpec{{
			Name:          "n1",
			Class:         "c1",
			Hostname:      "n1",
			IP:            "10.0.0.1",
			K3sRole:       "agent",
			SSHUser:       "atlas",
			SSHAuthorized: []string{"ssh-ed25519 AAA"},
		}},
	}
	out := filepath.Join(dir, "output.img")
	if err := BuildImageFile(context.Background(), inv, "n1", filepath.Join(dir, "cache"), out); err != nil {
		t.Fatalf("BuildImageFile: %v", err)
	}
	if _, err := os.Stat(out); err != nil {
		t.Fatalf("expected output image: %v", err)
	}
}

func TestMaybeInjectNoopsWhenEnvUnset(t *testing.T) {
	if err := maybeInject(&inventory.Inventory{}, "n1"); err != nil {
		t.Fatalf("maybeInject without env: %v", err)
	}
}

func TestChecksumFromInventoryAndCacheName(t *testing.T) {
	inv := &inventory.Inventory{
		Classes: []inventory.NodeClass{{Name: "c1", Checksum: "sha256:deadbeef"}},
		Nodes:   []inventory.NodeSpec{{Name: "n1", Class: "c1"}},
	}
	if got := checksumFromInventory(inv, "n1"); got != "sha256:deadbeef" {
		t.Fatalf("checksumFromInventory = %q", got)
	}
	if got := cacheName("/tmp/archive/base.img.xz"); got != "base.img" {
		t.Fatalf("cacheName = %q", got)
	}
}

func fakeRootfsTools(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	write := func(name, body string) {
		path := filepath.Join(dir, name)
		if err := os.WriteFile(path, []byte("#!/usr/bin/env bash\nset -eu\n"+body+"\n"), 0o755); err != nil {
			t.Fatalf("write %s: %v", name, err)
		}
	}
	write("sfdisk", `cat <<'JSON'
{"partitiontable":{"sectorsize":512,"partitions":[{"start":3,"size":1,"type":"ef"},{"start":1,"size":2,"type":"83"}]}}
JSON`)
	write("debugfs", `if [[ "${1:-}" == "-w" ]]; then
  cp "${3:-}" "${4:-}.commands"
  exit 0
fi
if [[ "${1:-}" == "-R" ]]; then
  state="${3:-}.commands"
  set -- $2
  case "${1:-}" in
    stat)
      mode="$(awk -v path="${2:-}" '$1=="sif" && $2==path {print $4}' "${state}" | tail -n1)"
      mode="${mode: -4}"
      printf 'Mode:  %s\n' "${mode}"
      exit 0
      ;;
    dump)
      local_path="$(awk -v path="${2:-}" '$1=="write" && $3==path {print $2}' "${state}" | tail -n1)"
      cat "${local_path}" > "${3:-}"
      exit 0
      ;;
  esac
fi
exit 0`)
	return dir
}

func fakeMountTools(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	write := func(name, body string) {
		path := filepath.Join(dir, name)
		if err := os.WriteFile(path, []byte("#!/usr/bin/env bash\nset -eu\n"+body+"\n"), 0o755); err != nil {
			t.Fatalf("write %s: %v", name, err)
		}
	}
	write("losetup", `if [[ "${1:-}" == "-Pf" && "${2:-}" == "--show" ]]; then
  printf '/dev/loop9\n'
  exit 0
fi
exit 0`)
	write("mount", `exit 0`)
	write("umount", `exit 0`)
	return dir
}

func imageChecksum(t *testing.T, path string) string {
	t.Helper()
	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	sum := sha256.Sum256(data)
	return "sha256:" + hex.EncodeToString(sum[:])
}