src/pkg/oempreloader/oem_test.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package oempreloader

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
)

const srcImage = "testdata/disk.img"

func TestExtend(t *testing.T) {
	diskSizeGB := 1
	oemFSSize := "10M"
	targetOEMPartitionSector := uint64(40960)
	tmpDir := t.TempDir()
	outImage := filepath.Join(tmpDir, "out.img")
	if err := Extend(srcImage, outImage, oemFSSize, diskSizeGB); err != nil {
		t.Fatalf("Failed to run Extend: %v", err)
	}
	size, err := partutil.ReadPartitionSize(outImage, OEMPartitionNum)
	if err != nil {
		t.Fatalf("Failed to read partition size: %v", err)
	}
	if size != targetOEMPartitionSector {
		t.Fatalf("Wrong partition size, expected: %d, actual: %d ", targetOEMPartitionSector, size)
	}
}

func TestPreloadDir(t *testing.T) {
	targetImage := "targetImage"
	if err := copyFile(srcImage, targetImage); err != nil {
		t.Fatalf("Failed to setup test image: %v", err)
	}
	defer os.Remove(targetImage)

	srcDir := "testdata/preload_dir"
	oemFSSize := "1M"
	partitionSize := int64(2097152)
	oemFile, err := PreloadDir(srcDir, oemFSSize, partitionSize)
	if err != nil {
		t.Fatalf("Failed to run PreloadDir: %v", err)
	}

	if err := WritePartitionFileToImage(oemFile, targetImage, OEMPartitionNum); err != nil {
		t.Fatalf("Failed to write oem to image: %v", err)
	}

	loopDev, err := setupLoopDevice(targetImage)
	if err != nil {
		t.Fatalf("Failed to setup loop device: %v", err)
	}
	defer detachLoopDevice(loopDev)
	partitionDevice := loopDev + "p8"
	tmpDir := t.TempDir()
	mountCmd := exec.Command("sudo", "mount", partitionDevice, tmpDir)
	if err := mountCmd.Run(); err != nil {
		t.Fatalf("Failed to mount partition: %v", err)
	}
	defer exec.Command("sudo", "umount", tmpDir).Run()

	data, err := os.ReadFile(tmpDir + "/a.txt")
	if err != nil {
		t.Fatalf("Could not read a.txt: %v", err)
	}
	expectedString := "1234567890"
	if string(data) != expectedString {
		t.Fatalf("Wrong data. Want: %q, got: %q", expectedString, string(data))
	}
}

func copyFile(src, dst string) error {
	sourceFile, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("Failed to open source image: %v", err)
	}
	defer sourceFile.Close()

	destFile, err := os.Create(dst)
	if err != nil {
		return fmt.Errorf("Failed to create test image: %v", err)
	}
	defer destFile.Close()
	_, err = io.Copy(destFile, sourceFile)
	return nil
}

func setupLoopDevice(loopDev string) (string, error) {
	cmd := exec.Command("sudo", "losetup", "-f", "-P", "--show", loopDev)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("failed to setup loop device: %v, output: %s", err, string(output))
	}
	// Output usually contains a newline, e.g., "/dev/loop0\n"
	loopDevice := strings.TrimSpace(string(output))
	if loopDevice == "" {
		return "", fmt.Errorf("losetup returned an empty device name")
	}
	return loopDevice, nil
}

func detachLoopDevice(device string) error {
	return exec.Command("sudo", "losetup", "-d", device).Run()
}