go-kv/db_test.go

package go_kv

import (
	"fmt"
	"os"
	"path/filepath"
	"testing"
)

func TestDBOpenClose(t *testing.T) {
	db, err := Open(t.TempDir(), nil)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	if err := db.Put([]byte("k"), []byte("v")); err != nil {
		t.Fatalf("Put: %v", err)
	}
	if err := db.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
}

func TestDBPutGet(t *testing.T) {
	db := openTestDB(t, t.TempDir())

	if err := db.Put([]byte("k1"), []byte("v1")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	got := db.Get([]byte("k1"))
	if !got.Found {
		t.Fatal("Get(k1) not found")
	}
	if string(got.Value) != "v1" {
		t.Fatalf("Get(k1) value = %q, want %q", got.Value, "v1")
	}
}

func TestDBDelete(t *testing.T) {
	db := openTestDB(t, t.TempDir())

	if err := db.Put([]byte("k"), []byte("v")); err != nil {
		t.Fatalf("Put: %v", err)
	}
	if err := db.Delete([]byte("k")); err != nil {
		t.Fatalf("Delete: %v", err)
	}

	got := db.Get([]byte("k"))
	if got.Found {
		t.Fatalf("Get(k) found deleted key with value %q", got.Value)
	}
}

func TestDBRecovery(t *testing.T) {
	dir := t.TempDir()
	db := openTestDB(t, dir)

	entries := map[string]string{
		"k1": "v1",
		"k2": "v2",
		"k3": "v3",
	}
	for key, value := range entries {
		if err := db.Put([]byte(key), []byte(value)); err != nil {
			t.Fatalf("Put(%q): %v", key, err)
		}
	}
	closeTestDB(t, db)

	reopened := openTestDB(t, dir)
	for key, want := range entries {
		got := reopened.Get([]byte(key))
		if !got.Found {
			t.Fatalf("Get(%q) not found after recovery", key)
		}
		if string(got.Value) != want {
			t.Fatalf("Get(%q) value = %q, want %q", key, got.Value, want)
		}
	}
}

func TestDBRecoveryAfterCorruption(t *testing.T) {
	dir := t.TempDir()
	db := openTestDB(t, dir)

	if err := db.Put([]byte("good1"), []byte("v1")); err != nil {
		t.Fatalf("Put good1: %v", err)
	}
	if err := db.Put([]byte("good2"), []byte("v2")); err != nil {
		t.Fatalf("Put good2: %v", err)
	}
	closeTestDB(t, db)

	appendWalTail(t, dir, []byte{0xde, 0xad, 0xbe, 0xef})

	reopened := openTestDB(t, dir)
	for key, want := range map[string]string{"good1": "v1", "good2": "v2"} {
		got := reopened.Get([]byte(key))
		if !got.Found {
			t.Fatalf("Get(%q) not found after corrupted-tail recovery", key)
		}
		if string(got.Value) != want {
			t.Fatalf("Get(%q) value = %q, want %q", key, got.Value, want)
		}
	}
}

func TestDBMultiplePuts(t *testing.T) {
	dir := t.TempDir()
	db := openTestDB(t, dir)

	for i := range 50 {
		key := fmt.Appendf(nil, "k%d", i)
		value := fmt.Appendf(nil, "v%d", i)
		if err := db.Put(key, value); err != nil {
			t.Fatalf("Put(%q): %v", key, err)
		}
	}
	closeTestDB(t, db)

	reopened := openTestDB(t, dir)
	for i := range 50 {
		key := fmt.Appendf(nil, "k%d", i)
		want := fmt.Sprintf("v%d", i)
		got := reopened.Get(key)
		if !got.Found {
			t.Fatalf("Get(%q) not found after recovery", key)
		}
		if string(got.Value) != want {
			t.Fatalf("Get(%q) value = %q, want %q", key, got.Value, want)
		}
	}
}

func openTestDB(t *testing.T, dir string) *DB {
	t.Helper()

	db, err := Open(dir, nil)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	t.Cleanup(func() {
		if !db.closed.Load() {
			if err := db.Close(); err != nil {
				t.Fatalf("Close cleanup: %v", err)
			}
		}
	})
	return db
}

func closeTestDB(t *testing.T, db *DB) {
	t.Helper()
	if err := db.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
}

func appendWalTail(t *testing.T, dir string, tail []byte) {
	t.Helper()

	path := filepath.Join(dir, "segment-0.wal")
	f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0)
	if err != nil {
		t.Fatalf("OpenFile(%s): %v", path, err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			t.Fatalf("Close corrupted WAL file: %v", err)
		}
	}()
	if _, err := f.Write(tail); err != nil {
		t.Fatalf("Write corrupted WAL tail: %v", err)
	}
}