metis/pkg/image/rootfs_progress_test.go

package image

import (
	"bytes"
	"io"
	"os"
	"path/filepath"
	"testing"

	"metis/pkg/inject"
)

type shortWriteOnly struct{}

func (shortWriteOnly) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	return len(p) - 1, nil
}

type failAfterFirstRead struct{ done bool }

func (r *failAfterFirstRead) Read(p []byte) (int, error) {
	if r.done {
		return 0, io.ErrUnexpectedEOF
	}
	r.done = true
	copy(p, []byte("ab"))
	return 2, nil
}

func TestInjectRootFSWithFakesReportsReplacingByteProgress(t *testing.T) {
	scripts := fakeRootfsCommands(t, true)
	t.Setenv("PATH", scripts+string(os.PathListSeparator)+os.Getenv("PATH"))

	imagePath := filepath.Join(t.TempDir(), "image.img")
	if err := os.WriteFile(imagePath, make([]byte, 4096), 0o644); err != nil {
		t.Fatal(err)
	}
	files := []inject.FileSpec{{Path: "etc/metis/firstboot.env", Content: []byte("METIS_HOSTNAME='titan-13'\n"), Mode: 0o600, RootFS: true}}
	var replacing []RootFSProgressUpdate
	if err := InjectRootFSWithProgress(imagePath, files, func(update RootFSProgressUpdate) {
		if update.Step == RootFSProgressReplacing && update.TotalBytes > 0 {
			replacing = append(replacing, update)
		}
	}); err != nil {
		t.Fatalf("InjectRootFSWithProgress: %v", err)
	}
	if len(replacing) == 0 {
		t.Fatal("expected replacing progress updates")
	}
	last := replacing[len(replacing)-1]
	if last.WrittenBytes != last.TotalBytes || last.TotalBytes != 1024 {
		t.Fatalf("unexpected replacing progress: %+v", last)
	}
}

func TestMkdirTempNearPathPrefersImageAdjacentScratch(t *testing.T) {
	imagePath := filepath.Join(t.TempDir(), "build", "node.img")
	if err := os.MkdirAll(filepath.Dir(imagePath), 0o755); err != nil {
		t.Fatal(err)
	}
	workDir, err := mkdirTempNearPath(imagePath, "metis-rootfs-")
	if err != nil {
		t.Fatalf("mkdirTempNearPath: %v", err)
	}
	defer os.RemoveAll(workDir)
	if got, want := filepath.Base(filepath.Dir(workDir)), ".metis-tmp"; got != want {
		t.Fatalf("temp parent = %q want %q", got, want)
	}
}

func TestCopyWithProgressReportsByteCounts(t *testing.T) {
	var dst bytes.Buffer
	var updates []RootFSProgressUpdate
	src := bytes.NewReader([]byte("hello world"))
	written, err := copyWithProgress(&dst, src, int64(src.Len()), func(written, total int64) {
		updates = append(updates, RootFSProgressUpdate{Step: RootFSProgressReplacing, WrittenBytes: written, TotalBytes: total})
	})
	if err != nil {
		t.Fatalf("copyWithProgress: %v", err)
	}
	if written != int64(len("hello world")) {
		t.Fatalf("written = %d", written)
	}
	if dst.String() != "hello world" {
		t.Fatalf("unexpected dst content: %q", dst.String())
	}
	if len(updates) == 0 || updates[len(updates)-1].WrittenBytes != updates[len(updates)-1].TotalBytes {
		t.Fatalf("expected final byte-complete update, got %#v", updates)
	}
}

func TestReplacePartitionWritesDataAndReportsProgress(t *testing.T) {
	dir := t.TempDir()
	imagePath := filepath.Join(dir, "image.img")
	rootPath := filepath.Join(dir, "root.ext4")
	imageBytes := make([]byte, 1024)
	rootBytes := bytes.Repeat([]byte("r"), 512)
	if err := os.WriteFile(imagePath, imageBytes, 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(rootPath, rootBytes, 0o644); err != nil {
		t.Fatal(err)
	}
	var updates []RootFSProgressUpdate
	if err := replacePartition(imagePath, rootPath, partitionTablePart{Start: 1, Size: 1}, 512, func(written, total int64) {
		updates = append(updates, RootFSProgressUpdate{Step: RootFSProgressReplacing, WrittenBytes: written, TotalBytes: total})
	}); err != nil {
		t.Fatalf("replacePartition: %v", err)
	}
	got, err := os.ReadFile(imagePath)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got[512:], rootBytes) {
		t.Fatalf("partition bytes were not replaced: %q", got[512:])
	}
	if len(updates) == 0 || updates[len(updates)-1].WrittenBytes != 512 || updates[len(updates)-1].TotalBytes != 512 {
		t.Fatalf("expected final replace progress update, got %#v", updates)
	}
}

func TestMkdirTempNearPathHonorsOverrideAndFallback(t *testing.T) {
	t.Run("override", func(t *testing.T) {
		override := filepath.Join(t.TempDir(), "override")
		t.Setenv("METIS_ROOTFS_TMP_DIR", override)
		workDir, err := mkdirTempNearPath(filepath.Join(t.TempDir(), "build", "node.img"), "metis-rootfs-")
		if err != nil {
			t.Fatalf("mkdirTempNearPath override: %v", err)
		}
		defer os.RemoveAll(workDir)
		if got := filepath.Dir(workDir); got != override {
			t.Fatalf("override parent = %q want %q", got, override)
		}
	})
	t.Run("fallback tempdir", func(t *testing.T) {
		workDir, err := mkdirTempNearPath("", "metis-rootfs-")
		if err != nil {
			t.Fatalf("mkdirTempNearPath fallback: %v", err)
		}
		defer os.RemoveAll(workDir)
		if workDir == "" {
			t.Fatal("expected fallback temp dir path")
		}
	})
}

func TestCopyWithProgressPropagatesReaderAndWriterErrors(t *testing.T) {
	t.Run("short write", func(t *testing.T) {
		if _, err := copyWithProgress(shortWriteOnly{}, bytes.NewReader([]byte("abc")), 3, nil); err != io.ErrShortWrite {
			t.Fatalf("expected short write, got %v", err)
		}
	})
	t.Run("reader error", func(t *testing.T) {
		if _, err := copyWithProgress(&bytes.Buffer{}, &failAfterFirstRead{}, 4, nil); err != io.ErrUnexpectedEOF {
			t.Fatalf("expected reader error, got %v", err)
		}
	})
}