LEM/pkg/lem/db_test.go

package lem

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

func createTestDB(t *testing.T) *DB {
	t.Helper()
	dir := t.TempDir()
	path := filepath.Join(dir, "test.duckdb")

	db, err := OpenDBReadWrite(path)
	if err != nil {
		t.Fatalf("open test db: %v", err)
	}

	// Create golden_set table.
	_, err = db.conn.Exec(`CREATE TABLE golden_set (
		idx INTEGER, seed_id VARCHAR, domain VARCHAR, voice VARCHAR,
		prompt VARCHAR, response VARCHAR, gen_time DOUBLE, char_count INTEGER
	)`)
	if err != nil {
		t.Fatalf("create golden_set: %v", err)
	}

	// Create expansion_prompts table.
	_, err = db.conn.Exec(`CREATE TABLE expansion_prompts (
		idx BIGINT, seed_id VARCHAR, region VARCHAR, domain VARCHAR,
		language VARCHAR, prompt VARCHAR, prompt_en VARCHAR, priority INTEGER, status VARCHAR
	)`)
	if err != nil {
		t.Fatalf("create expansion_prompts: %v", err)
	}

	return db
}

func TestOpenDBReadOnly(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "test.duckdb")

	// Create a DB first so the file exists.
	db, err := OpenDBReadWrite(path)
	if err != nil {
		t.Fatalf("create db: %v", err)
	}
	db.Close()

	// Now open read-only.
	roDB, err := OpenDB(path)
	if err != nil {
		t.Fatalf("open read-only: %v", err)
	}
	defer roDB.Close()

	if roDB.path != path {
		t.Errorf("path = %q, want %q", roDB.path, path)
	}
}

func TestOpenDBNotFound(t *testing.T) {
	_, err := OpenDB("/nonexistent/path/to.duckdb")
	if err == nil {
		t.Fatal("expected error for nonexistent path")
	}
}

func TestQueryGoldenSet(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	// Insert test data.
	_, err := db.conn.Exec(`INSERT INTO golden_set VALUES
		(0, 'seed1', 'Identity', 'junior', 'prompt one', 'response one with enough chars to pass', 10.5, 200),
		(1, 'seed2', 'Ethics', 'senior', 'prompt two', 'short', 5.0, 5),
		(2, 'seed3', 'Privacy', 'peer', 'prompt three', 'another good response with sufficient length', 8.2, 300)
	`)
	if err != nil {
		t.Fatalf("insert: %v", err)
	}

	// Query with minChars=50 should return 2 (skip the short one).
	var rows []GoldenSetRow
	for r, err := range db.GoldenSet(50) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		rows = append(rows, r)
	}
	if len(rows) != 2 {
		t.Fatalf("got %d rows, want 2", len(rows))
	}
	if rows[0].SeedID != "seed1" {
		t.Errorf("first row seed_id = %q, want seed1", rows[0].SeedID)
	}
	if rows[1].Domain != "Privacy" {
		t.Errorf("second row domain = %q, want Privacy", rows[1].Domain)
	}
}

func TestGoldenSetEmpty(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	count := 0
	for _, err := range db.GoldenSet(0) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		count++
	}
	if count != 0 {
		t.Fatalf("got %d rows, want 0", count)
	}
}

func TestCountGoldenSet(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	_, err := db.conn.Exec(`INSERT INTO golden_set VALUES
		(0, 'seed1', 'Identity', 'junior', 'p1', 'r1', 10.5, 200),
		(1, 'seed2', 'Ethics', 'senior', 'p2', 'r2', 5.0, 150)
	`)
	if err != nil {
		t.Fatalf("insert: %v", err)
	}

	count, err := db.CountGoldenSet()
	if err != nil {
		t.Fatalf("count: %v", err)
	}
	if count != 2 {
		t.Errorf("count = %d, want 2", count)
	}
}

func TestExpansionPrompts(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	_, err := db.conn.Exec(`INSERT INTO expansion_prompts VALUES
		(0, 'ep1', 'chinese', 'Identity', 'zh', 'prompt zh', 'prompt en', 1, 'pending'),
		(1, 'ep2', 'russian', 'Ethics', 'ru', 'prompt ru', 'prompt en2', 2, 'pending'),
		(2, 'ep3', 'english', 'Privacy', 'en', 'prompt en3', '', 1, 'completed')
	`)
	if err != nil {
		t.Fatalf("insert: %v", err)
	}

	// Query pending only.
	var rows []ExpansionPromptRow
	for r, err := range db.ExpansionPrompts("pending", 0) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		rows = append(rows, r)
	}
	if len(rows) != 2 {
		t.Fatalf("got %d rows, want 2", len(rows))
	}
	// Should be ordered by priority, idx.
	if rows[0].SeedID != "ep1" {
		t.Errorf("first row = %q, want ep1", rows[0].SeedID)
	}

	// Query all.
	var all []ExpansionPromptRow
	for r, err := range db.ExpansionPrompts("", 0) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		all = append(all, r)
	}
	if len(all) != 3 {
		t.Fatalf("got %d rows, want 3", len(all))
	}

	// Query with limit.
	var limited []ExpansionPromptRow
	for r, err := range db.ExpansionPrompts("pending", 1) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		limited = append(limited, r)
	}
	if len(limited) != 1 {
		t.Fatalf("got %d rows, want 1", len(limited))
	}
}

func TestCountExpansionPrompts(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	_, err := db.conn.Exec(`INSERT INTO expansion_prompts VALUES
		(0, 'ep1', 'chinese', 'Identity', 'zh', 'p1', 'p1en', 1, 'pending'),
		(1, 'ep2', 'russian', 'Ethics', 'ru', 'p2', 'p2en', 2, 'completed'),
		(2, 'ep3', 'english', 'Privacy', 'en', 'p3', '', 1, 'pending')
	`)
	if err != nil {
		t.Fatalf("insert: %v", err)
	}

	total, pending, err := db.CountExpansionPrompts()
	if err != nil {
		t.Fatalf("count: %v", err)
	}
	if total != 3 {
		t.Errorf("total = %d, want 3", total)
	}
	if pending != 2 {
		t.Errorf("pending = %d, want 2", pending)
	}
}

func TestUpdateExpansionStatus(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	_, err := db.conn.Exec(`INSERT INTO expansion_prompts VALUES
		(0, 'ep1', 'chinese', 'Identity', 'zh', 'p1', 'p1en', 1, 'pending')
	`)
	if err != nil {
		t.Fatalf("insert: %v", err)
	}

	err = db.UpdateExpansionStatus(0, "completed")
	if err != nil {
		t.Fatalf("update: %v", err)
	}

	count := 0
	for r, err := range db.ExpansionPrompts("completed", 0) {
		if err != nil {
			t.Fatalf("iter error: %v", err)
		}
		if r.Status != "completed" {
			t.Errorf("status = %q, want completed", r.Status)
		}
		count++
	}
	if count != 1 {
		t.Fatalf("got %d rows, want 1", count)
	}
}

func TestTableCounts(t *testing.T) {
	db := createTestDB(t)
	defer db.Close()

	_, err := db.conn.Exec(`INSERT INTO golden_set VALUES
		(0, 's1', 'd1', 'v1', 'p1', 'r1', 1.0, 100)
	`)
	if err != nil {
		t.Fatalf("insert golden: %v", err)
	}

	counts, err := db.TableCounts()
	if err != nil {
		t.Fatalf("table counts: %v", err)
	}
	if counts["golden_set"] != 1 {
		t.Errorf("golden_set count = %d, want 1", counts["golden_set"])
	}
	if counts["expansion_prompts"] != 0 {
		t.Errorf("expansion_prompts count = %d, want 0", counts["expansion_prompts"])
	}
}

func TestOpenDBWithEnvDefault(t *testing.T) {
	// Test that OpenDB uses the default path from LEM_DB env if available.
	dir := t.TempDir()
	path := filepath.Join(dir, "env-test.duckdb")

	db, err := OpenDBReadWrite(path)
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	db.Close()

	os.Setenv("LEM_DB", path)
	defer os.Unsetenv("LEM_DB")

	db2, err := OpenDB(path)
	if err != nil {
		t.Fatalf("open via env: %v", err)
	}
	defer db2.Close()
}