kubesolo-os/update/pkg/image/image_test.go

package image

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
)

func TestCheckForUpdate(t *testing.T) {
	meta := UpdateMetadata{
		Version:         "1.2.0",
		VmlinuzURL:      "/vmlinuz",
		VmlinuzSHA256:   "abc123",
		InitramfsURL:    "/kubesolo-os.gz",
		InitramfsSHA256: "def456",
		ReleaseNotes:    "Bug fixes",
		ReleaseDate:     "2025-01-15",
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/latest.json" {
			http.NotFound(w, r)
			return
		}
		json.NewEncoder(w).Encode(meta)
	}))
	defer server.Close()

	client := NewClient(server.URL, "")
	got, err := client.CheckForUpdate()
	if err != nil {
		t.Fatal(err)
	}

	if got.Version != "1.2.0" {
		t.Errorf("expected version 1.2.0, got %s", got.Version)
	}
	if got.VmlinuzSHA256 != "abc123" {
		t.Errorf("expected vmlinuz sha abc123, got %s", got.VmlinuzSHA256)
	}
	if got.ReleaseNotes != "Bug fixes" {
		t.Errorf("expected release notes, got %s", got.ReleaseNotes)
	}
}

func TestCheckForUpdateMissingVersion(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(UpdateMetadata{})
	}))
	defer server.Close()

	client := NewClient(server.URL, "")
	_, err := client.CheckForUpdate()
	if err == nil {
		t.Fatal("expected error for missing version")
	}
}

func TestCheckForUpdateServerError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer server.Close()

	client := NewClient(server.URL, "")
	_, err := client.CheckForUpdate()
	if err == nil {
		t.Fatal("expected error for server error")
	}
}

func TestDownloadAndVerify(t *testing.T) {
	// Create test content
	vmlinuzContent := []byte("fake vmlinuz content for testing")
	initramfsContent := []byte("fake initramfs content for testing")

	vmlinuzHash := sha256.Sum256(vmlinuzContent)
	initramfsHash := sha256.Sum256(initramfsContent)

	meta := UpdateMetadata{
		Version:         "2.0.0",
		VmlinuzSHA256:   hex.EncodeToString(vmlinuzHash[:]),
		InitramfsSHA256: hex.EncodeToString(initramfsHash[:]),
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/latest.json":
			m := meta
			m.VmlinuzURL = "http://" + r.Host + "/vmlinuz"
			m.InitramfsURL = "http://" + r.Host + "/kubesolo-os.gz"
			json.NewEncoder(w).Encode(m)
		case "/vmlinuz":
			w.Write(vmlinuzContent)
		case "/kubesolo-os.gz":
			w.Write(initramfsContent)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	stageDir := filepath.Join(t.TempDir(), "stage")
	client := NewClient(server.URL, stageDir)
	defer client.Cleanup()

	// First get metadata
	gotMeta, err := client.CheckForUpdate()
	if err != nil {
		t.Fatal(err)
	}

	// Download
	staged, err := client.Download(gotMeta)
	if err != nil {
		t.Fatal(err)
	}

	if staged.Version != "2.0.0" {
		t.Errorf("expected version 2.0.0, got %s", staged.Version)
	}

	// Verify files exist
	if _, err := os.Stat(staged.VmlinuzPath); err != nil {
		t.Errorf("vmlinuz not found: %v", err)
	}
	if _, err := os.Stat(staged.InitramfsPath); err != nil {
		t.Errorf("initramfs not found: %v", err)
	}

	// Verify content
	data, _ := os.ReadFile(staged.VmlinuzPath)
	if string(data) != string(vmlinuzContent) {
		t.Error("vmlinuz content mismatch")
	}
}

func TestDownloadChecksumMismatch(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/vmlinuz":
			w.Write([]byte("actual content"))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	stageDir := filepath.Join(t.TempDir(), "stage")
	client := NewClient(server.URL, stageDir)

	meta := &UpdateMetadata{
		Version:       "1.0.0",
		VmlinuzURL:    server.URL + "/vmlinuz",
		VmlinuzSHA256: "wrong_checksum_value",
		InitramfsURL:  server.URL + "/initramfs",
	}

	_, err := client.Download(meta)
	if err == nil {
		t.Fatal("expected checksum mismatch error")
	}
}

func TestVerifyFile(t *testing.T) {
	content := []byte("test file content for verification")
	hash := sha256.Sum256(content)
	expected := hex.EncodeToString(hash[:])

	dir := t.TempDir()
	path := filepath.Join(dir, "testfile")
	if err := os.WriteFile(path, content, 0o644); err != nil {
		t.Fatal(err)
	}

	// Should pass with correct hash
	if err := VerifyFile(path, expected); err != nil {
		t.Errorf("expected verification to pass: %v", err)
	}

	// Should fail with wrong hash
	if err := VerifyFile(path, "deadbeef"); err == nil {
		t.Error("expected verification to fail with wrong hash")
	}
}

func TestVerifyFileNotFound(t *testing.T) {
	err := VerifyFile("/nonexistent/file", "abc123")
	if err == nil {
		t.Error("expected error for nonexistent file")
	}
}

func TestCleanup(t *testing.T) {
	stageDir := filepath.Join(t.TempDir(), "stage")
	os.MkdirAll(stageDir, 0o755)
	os.WriteFile(filepath.Join(stageDir, "test"), []byte("data"), 0o644)

	client := NewClient("http://unused", stageDir)
	if err := client.Cleanup(); err != nil {
		t.Fatal(err)
	}

	if _, err := os.Stat(stageDir); !os.IsNotExist(err) {
		t.Error("stage dir should be removed after cleanup")
	}
}

func TestUpdateMetadataJSON(t *testing.T) {
	meta := UpdateMetadata{
		Version:         "1.0.0",
		VmlinuzURL:      "https://example.com/vmlinuz",
		VmlinuzSHA256:   "abc",
		InitramfsURL:    "https://example.com/kubesolo-os.gz",
		InitramfsSHA256: "def",
		ReleaseNotes:    "Initial release",
		ReleaseDate:     "2025-01-01",
	}

	data, err := json.Marshal(meta)
	if err != nil {
		t.Fatal(err)
	}

	var decoded UpdateMetadata
	if err := json.Unmarshal(data, &decoded); err != nil {
		t.Fatal(err)
	}

	if decoded.Version != meta.Version {
		t.Errorf("version mismatch: %s != %s", decoded.Version, meta.Version)
	}
	if decoded.ReleaseDate != meta.ReleaseDate {
		t.Errorf("release date mismatch: %s != %s", decoded.ReleaseDate, meta.ReleaseDate)
	}
}