ExploreDNS/internal/receiver/store/store_test.go

package store

import (
	"context"
	"encoding/json"
	"path/filepath"
	"testing"
	"time"
)

func newTestStore(t *testing.T) *Store {
	t.Helper()
	s, err := OpenSQLite(filepath.Join(t.TempDir(), "receiver.db"))
	if err != nil {
		t.Fatalf("OpenSQLite: %v", err)
	}
	t.Cleanup(func() { s.Close() })
	return s
}

func startEvent(id string) StartEvent {
	return StartEvent{
		Event:     "start",
		ID:        id,
		Domain:    "example.com",
		QueryType: "A",
		AllRoots:  true,
		ClientIP:  "203.0.113.9",
		StartedAt: time.Date(2026, 7, 1, 10, 0, 0, 0, time.UTC),
	}
}

func completeEvent(id string) CompleteEvent {
	return CompleteEvent{
		Event:       "complete",
		ID:          id,
		Domain:      "example.com",
		QueryType:   "A",
		ClientIP:    "203.0.113.9",
		StartedAt:   time.Date(2026, 7, 1, 10, 0, 0, 0, time.UTC),
		DoneAt:      time.Date(2026, 7, 1, 10, 0, 3, 0, time.UTC),
		DurationMS:  3000,
		Status:      StatusComplete,
		ResultCount: 4,
		Summary:     json.RawMessage(`{"answers":[{"probability":1,"records":["example.com 300 IN A 192.0.2.1"]}]}`),
	}
}

// getRow fetches a single traversal by exact id via ListTraversals.
func getRow(t *testing.T, s *Store, id string) Traversal {
	t.Helper()
	rows, _, err := s.ListTraversals(context.Background(), ListFilter{}, 1000, 0)
	if err != nil {
		t.Fatalf("ListTraversals: %v", err)
	}
	for _, r := range rows {
		if r.ID == id {
			return r
		}
	}
	t.Fatalf("row %q not found", id)
	return Traversal{}
}

func TestRecordStartThenComplete(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	if err := s.RecordStart(ctx, startEvent("t1")); err != nil {
		t.Fatalf("RecordStart: %v", err)
	}
	row := getRow(t, s, "t1")
	if row.Status != StatusRunning {
		t.Errorf("Status = %q, want %q", row.Status, StatusRunning)
	}
	if !row.AllRoots {
		t.Error("AllRoots = false, want true")
	}
	if row.DoneAt != nil || row.DurationMS != nil || row.ResultCount != nil {
		t.Errorf("completion fields set before complete: %+v", row)
	}
	if !row.StartedAt.Equal(startEvent("t1").StartedAt) {
		t.Errorf("StartedAt = %v, want %v", row.StartedAt, startEvent("t1").StartedAt)
	}

	ev := completeEvent("t1")
	if err := s.RecordComplete(ctx, ev); err != nil {
		t.Fatalf("RecordComplete: %v", err)
	}
	row = getRow(t, s, "t1")
	if row.Status != StatusComplete {
		t.Errorf("Status = %q, want %q", row.Status, StatusComplete)
	}
	if row.DoneAt == nil || !row.DoneAt.Equal(ev.DoneAt) {
		t.Errorf("DoneAt = %v, want %v", row.DoneAt, ev.DoneAt)
	}
	if row.DurationMS == nil || *row.DurationMS != 3000 {
		t.Errorf("DurationMS = %v, want 3000", row.DurationMS)
	}
	if row.ResultCount == nil || *row.ResultCount != 4 {
		t.Errorf("ResultCount = %v, want 4", row.ResultCount)
	}
	if row.Summary != string(ev.Summary) {
		t.Errorf("Summary = %q, want %q", row.Summary, ev.Summary)
	}
	if !row.AllRoots {
		t.Error("AllRoots clobbered by complete event")
	}
}

func TestCompleteBeforeStart(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	if err := s.RecordComplete(ctx, completeEvent("t1")); err != nil {
		t.Fatalf("RecordComplete: %v", err)
	}
	row := getRow(t, s, "t1")
	if row.Status != StatusComplete {
		t.Fatalf("Status = %q, want %q", row.Status, StatusComplete)
	}

	// The delayed start must fill in start-only fields without touching
	// the terminal state.
	if err := s.RecordStart(ctx, startEvent("t1")); err != nil {
		t.Fatalf("RecordStart: %v", err)
	}
	row = getRow(t, s, "t1")
	if row.Status != StatusComplete {
		t.Errorf("Status = %q after late start, want %q", row.Status, StatusComplete)
	}
	if row.DoneAt == nil || row.DurationMS == nil {
		t.Errorf("completion fields lost after late start: %+v", row)
	}
	if !row.AllRoots {
		t.Error("AllRoots not filled in by late start")
	}
}

func TestDuplicateStartAfterComplete(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	if err := s.RecordStart(ctx, startEvent("t1")); err != nil {
		t.Fatalf("RecordStart: %v", err)
	}
	if err := s.RecordComplete(ctx, completeEvent("t1")); err != nil {
		t.Fatalf("RecordComplete: %v", err)
	}
	if err := s.RecordStart(ctx, startEvent("t1")); err != nil {
		t.Fatalf("duplicate RecordStart: %v", err)
	}

	row := getRow(t, s, "t1")
	if row.Status != StatusComplete {
		t.Errorf("Status = %q after duplicate start, want %q", row.Status, StatusComplete)
	}
	if row.DoneAt == nil {
		t.Error("DoneAt lost after duplicate start")
	}
	if row.Summary == "" {
		t.Error("Summary lost after duplicate start")
	}
}

func TestDuplicateCompleteIdempotent(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	ev := completeEvent("t1")
	ev.Status = StatusError
	ev.Error = "traversal timed out after 5m0s"
	for i := 0; i < 3; i++ {
		if err := s.RecordComplete(ctx, ev); err != nil {
			t.Fatalf("RecordComplete #%d: %v", i, err)
		}
	}

	rows, total, err := s.ListTraversals(ctx, ListFilter{}, 10, 0)
	if err != nil {
		t.Fatalf("ListTraversals: %v", err)
	}
	if total != 1 || len(rows) != 1 {
		t.Fatalf("total = %d, len = %d, want 1 row", total, len(rows))
	}
	if rows[0].Status != StatusError || rows[0].Error != ev.Error {
		t.Errorf("row = %q/%q, want %q/%q", rows[0].Status, rows[0].Error, StatusError, ev.Error)
	}
}

// seedRows inserts a deterministic mixed dataset for list/aggregate tests.
func seedRows(t *testing.T, s *Store) {
	t.Helper()
	ctx := context.Background()
	base := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
	rows := []struct {
		id, domain, qtype, ip, status string
		day                           int
		durMS                         int64
	}{
		{"a1", "example.com", "A", "203.0.113.1", StatusComplete, 0, 100},
		{"a2", "example.com", "AAAA", "203.0.113.1", StatusComplete, 0, 200},
		{"a3", "sub.example.com", "A", "203.0.113.2", StatusError, 1, 300},
		{"a4", "other.net", "MX", "203.0.113.3", StatusComplete, 1, 400},
		{"a5", "other.net", "A", "203.0.113.1", StatusComplete, 2, 500},
		{"a6", "under_score.org", "A", "203.0.113.4", StatusComplete, 2, 600},
	}
	for i, r := range rows {
		started := base.AddDate(0, 0, r.day).Add(time.Duration(i) * time.Minute)
		ev := CompleteEvent{
			ID: r.id, Domain: r.domain, QueryType: r.qtype, ClientIP: r.ip,
			StartedAt: started, DoneAt: started.Add(time.Duration(r.durMS) * time.Millisecond),
			DurationMS: r.durMS, Status: r.status, ResultCount: 1,
		}
		if r.status == StatusError {
			ev.Error = "lookup failed"
		}
		if err := s.RecordComplete(ctx, ev); err != nil {
			t.Fatalf("seed %s: %v", r.id, err)
		}
	}
	// One still-running traversal on day 2.
	if err := s.RecordStart(ctx, StartEvent{
		ID: "a7", Domain: "running.io", QueryType: "A", ClientIP: "203.0.113.5",
		StartedAt: base.AddDate(0, 0, 2).Add(time.Hour),
	}); err != nil {
		t.Fatalf("seed a7: %v", err)
	}
}

func TestListTraversalsFiltersAndPagination(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)
	ctx := context.Background()

	rows, total, err := s.ListTraversals(ctx, ListFilter{}, 3, 0)
	if err != nil {
		t.Fatalf("ListTraversals: %v", err)
	}
	if total != 7 || len(rows) != 3 {
		t.Fatalf("total = %d, len = %d, want 7 and 3", total, len(rows))
	}
	// Newest first: a7 (day2+1h) then a6, a5.
	if rows[0].ID != "a7" || rows[1].ID != "a6" || rows[2].ID != "a5" {
		t.Errorf("page 1 order = %s,%s,%s, want a7,a6,a5", rows[0].ID, rows[1].ID, rows[2].ID)
	}
	rows, _, err = s.ListTraversals(ctx, ListFilter{}, 3, 3)
	if err != nil {
		t.Fatalf("ListTraversals page 2: %v", err)
	}
	if rows[0].ID != "a4" || rows[1].ID != "a3" || rows[2].ID != "a2" {
		t.Errorf("page 2 order = %s,%s,%s, want a4,a3,a2", rows[0].ID, rows[1].ID, rows[2].ID)
	}

	// Domain substring matches example.com and sub.example.com.
	rows, total, err = s.ListTraversals(ctx, ListFilter{Domain: "example"}, 0, 0)
	if err != nil {
		t.Fatalf("domain filter: %v", err)
	}
	if total != 3 {
		t.Errorf("domain filter total = %d, want 3", total)
	}
	for _, r := range rows {
		if r.Domain != "example.com" && r.Domain != "sub.example.com" {
			t.Errorf("domain filter matched %q", r.Domain)
		}
	}

	// LIKE metacharacters in the filter must be literal, not wildcards.
	_, total, err = s.ListTraversals(ctx, ListFilter{Domain: "%"}, 0, 0)
	if err != nil {
		t.Fatalf("percent filter: %v", err)
	}
	if total != 0 {
		t.Errorf("%% filter total = %d, want 0", total)
	}
	_, total, err = s.ListTraversals(ctx, ListFilter{Domain: "under_score"}, 0, 0)
	if err != nil {
		t.Fatalf("underscore filter: %v", err)
	}
	if total != 1 {
		t.Errorf("under_score filter total = %d, want 1", total)
	}

	_, total, err = s.ListTraversals(ctx, ListFilter{Status: StatusError}, 0, 0)
	if err != nil {
		t.Fatalf("status filter: %v", err)
	}
	if total != 1 {
		t.Errorf("status filter total = %d, want 1", total)
	}

	base := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
	rows, total, err = s.ListTraversals(ctx, ListFilter{
		From: base.AddDate(0, 0, 1),
		To:   base.AddDate(0, 0, 2).Add(-time.Second),
	}, 0, 0)
	if err != nil {
		t.Fatalf("time filter: %v", err)
	}
	if total != 2 {
		t.Fatalf("time filter total = %d, want 2 (got %+v)", total, rows)
	}
	if rows[0].ID != "a4" || rows[1].ID != "a3" {
		t.Errorf("time filter = %s,%s, want a4,a3", rows[0].ID, rows[1].ID)
	}
}

func TestStatsPerDay(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	for i, spec := range []struct {
		daysAgo int
		status  string
	}{
		{0, StatusComplete}, {0, StatusError}, {1, StatusComplete}, {5, StatusComplete},
	} {
		ev := completeEvent(string(rune('a' + i)))
		ev.StartedAt = now.AddDate(0, 0, -spec.daysAgo)
		ev.DoneAt = ev.StartedAt.Add(time.Second)
		ev.Status = spec.status
		if err := s.RecordComplete(ctx, ev); err != nil {
			t.Fatalf("seed: %v", err)
		}
	}

	days, err := s.StatsPerDay(ctx, 3)
	if err != nil {
		t.Fatalf("StatsPerDay: %v", err)
	}
	if len(days) != 2 {
		t.Fatalf("len = %d, want 2 (%+v)", len(days), days)
	}
	yesterday := now.AddDate(0, 0, -1).Format("2006-01-02")
	today := now.Format("2006-01-02")
	if days[0].Day != yesterday || days[0].Total != 1 || days[0].Errors != 0 {
		t.Errorf("day[0] = %+v, want %s total 1 errors 0", days[0], yesterday)
	}
	if days[1].Day != today || days[1].Total != 2 || days[1].Errors != 1 {
		t.Errorf("day[1] = %+v, want %s total 2 errors 1", days[1], today)
	}
}

func TestTopDomains(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)

	top, err := s.TopDomains(context.Background(), time.Time{}, 2)
	if err != nil {
		t.Fatalf("TopDomains: %v", err)
	}
	want := []NameCount{{"example.com", 2}, {"other.net", 2}}
	if len(top) != 2 || top[0] != want[0] || top[1] != want[1] {
		t.Errorf("TopDomains = %+v, want %+v", top, want)
	}

	// Window excludes day 0 rows (a1, a2), so example.com drops out and
	// other.net (a4, a5) leads.
	since := time.Date(2026, 7, 2, 0, 0, 0, 0, time.UTC)
	top, err = s.TopDomains(context.Background(), since, 10)
	if err != nil {
		t.Fatalf("TopDomains since: %v", err)
	}
	if len(top) != 4 || top[0] != (NameCount{"other.net", 2}) {
		t.Errorf("TopDomains since = %+v, want other.net x2 leading 4 domains", top)
	}
}

func TestQueryTypeCounts(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)

	counts, err := s.QueryTypeCounts(context.Background(), time.Time{})
	if err != nil {
		t.Fatalf("QueryTypeCounts: %v", err)
	}
	want := []NameCount{{"A", 5}, {"AAAA", 1}, {"MX", 1}}
	if len(counts) != 3 || counts[0] != want[0] || counts[1] != want[1] || counts[2] != want[2] {
		t.Errorf("QueryTypeCounts = %+v, want %+v", counts, want)
	}
}

func TestStatusCounts(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)

	counts, err := s.StatusCounts(context.Background(), time.Time{})
	if err != nil {
		t.Fatalf("StatusCounts: %v", err)
	}
	want := []NameCount{{StatusComplete, 5}, {StatusError, 1}, {StatusRunning, 1}}
	if len(counts) != 3 || counts[0] != want[0] || counts[1] != want[1] || counts[2] != want[2] {
		t.Errorf("StatusCounts = %+v, want %+v", counts, want)
	}
}

func TestTopClientIPs(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)

	top, err := s.TopClientIPs(context.Background(), time.Time{}, 1)
	if err != nil {
		t.Fatalf("TopClientIPs: %v", err)
	}
	if len(top) != 1 || top[0].Name != "203.0.113.1" || top[0].Count != 3 {
		t.Errorf("TopClientIPs = %+v, want 203.0.113.1 x3", top)
	}
}

func TestDurations(t *testing.T) {
	s := newTestStore(t)
	seedRows(t, s)

	// Durations 100..600; the running row has none and is excluded.
	stats, err := s.Durations(context.Background(), time.Time{})
	if err != nil {
		t.Fatalf("Durations: %v", err)
	}
	if stats.Count != 6 {
		t.Fatalf("Count = %d, want 6", stats.Count)
	}
	if stats.AvgMS != 350 {
		t.Errorf("AvgMS = %v, want 350", stats.AvgMS)
	}
	if stats.P50MS != 300 {
		t.Errorf("P50MS = %d, want 300", stats.P50MS)
	}
	if stats.P95MS != 600 {
		t.Errorf("P95MS = %d, want 600", stats.P95MS)
	}

	// Empty window.
	stats, err = s.Durations(context.Background(), time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC))
	if err != nil {
		t.Fatalf("Durations empty: %v", err)
	}
	if stats.Count != 0 || stats.AvgMS != 0 || stats.P50MS != 0 || stats.P95MS != 0 {
		t.Errorf("empty Durations = %+v, want zeros", stats)
	}
}

func TestTotals(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	tot, err := s.Totals(ctx, now)
	if err != nil {
		t.Fatalf("Totals empty: %v", err)
	}
	if tot != (Totals{}) {
		t.Errorf("empty Totals = %+v, want zeros", tot)
	}

	for _, spec := range []struct {
		id, domain, ip, status string
		ago                    time.Duration
	}{
		{"t1", "a.com", "203.0.113.1", StatusComplete, time.Hour},
		{"t2", "a.com", "203.0.113.2", StatusError, 2 * time.Hour},
		{"t3", "b.net", "203.0.113.1", StatusComplete, 48 * time.Hour},
		{"t4", "c.org", "203.0.113.3", StatusComplete, 10 * 24 * time.Hour},
	} {
		ev := completeEvent(spec.id)
		ev.Domain = spec.domain
		ev.ClientIP = spec.ip
		ev.Status = spec.status
		ev.StartedAt = now.Add(-spec.ago)
		ev.DoneAt = ev.StartedAt.Add(time.Second)
		if err := s.RecordComplete(ctx, ev); err != nil {
			t.Fatalf("seed %s: %v", spec.id, err)
		}
	}

	tot, err = s.Totals(ctx, now)
	if err != nil {
		t.Fatalf("Totals: %v", err)
	}
	want := Totals{AllTime: 4, Last24h: 2, Last7d: 3, DistinctDomains: 3, DistinctClients: 3, Errors: 1}
	if tot != want {
		t.Errorf("Totals = %+v, want %+v", tot, want)
	}
}

func TestOpenSQLiteCreatesParentDirs(t *testing.T) {
	path := filepath.Join(t.TempDir(), "nested", "dir", "receiver.db")
	s, err := OpenSQLite(path)
	if err != nil {
		t.Fatalf("OpenSQLite: %v", err)
	}
	defer s.Close()
	if s.Backend() != "sqlite" {
		t.Errorf("Backend = %q, want sqlite", s.Backend())
	}

	// Reopening must not fail on the existing schema and must keep the
	// stamped version.
	s.Close()
	s2, err := OpenSQLite(path)
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer s2.Close()
	var v int
	if err := s2.db.QueryRow(`SELECT version FROM schema_version`).Scan(&v); err != nil {
		t.Fatalf("schema_version: %v", err)
	}
	if v != schemaVersion {
		t.Errorf("schema version = %d, want %d", v, schemaVersion)
	}
}

func TestRedactMySQLDSN(t *testing.T) {
	got := RedactMySQLDSN("user:s3cret@tcp(db.example.com:3306)/exploredns")
	if got != "user:xxxxx@tcp(db.example.com:3306)/exploredns" {
		t.Errorf("RedactMySQLDSN = %q", got)
	}
	if got := RedactMySQLDSN("::::"); got != "(unparsable DSN)" {
		t.Errorf("RedactMySQLDSN(bad) = %q", got)
	}
}

// TestMySQLDialectSQL pins the MySQL-side SQL text, which unit tests cannot
// execute without a server.
func TestMySQLDialectSQL(t *testing.T) {
	d := mysqlDialect{}
	got := d.upsert("INSERT INTO traversals (id, domain) VALUES (?, ?)", []string{"domain", "last_seen"})
	want := "INSERT INTO traversals (id, domain) VALUES (?, ?)" +
		" ON DUPLICATE KEY UPDATE domain=VALUES(domain), last_seen=VALUES(last_seen)"
	if got != want {
		t.Errorf("upsert = %q, want %q", got, want)
	}
	pred, arg := d.likeContains("domain", `50%_o\ff`)
	if pred != `domain LIKE ? ESCAPE '\\'` {
		t.Errorf("likeContains predicate = %q", pred)
	}
	if arg != `%50\%\_o\\ff%` {
		t.Errorf("likeContains arg = %q", arg)
	}
	if d.datetimeType() != "DATETIME(6)" {
		t.Errorf("datetimeType = %q", d.datetimeType())
	}
}