oubliette/internal/storage/store_test.go

package storage

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

// storeFactory returns a clean Store and a cleanup function.
type storeFactory func(t *testing.T) Store

func testStores(t *testing.T, f func(t *testing.T, newStore storeFactory)) {
	t.Helper()
	t.Run("SQLite", func(t *testing.T) {
		f(t, func(t *testing.T) Store {
			t.Helper()
			dbPath := filepath.Join(t.TempDir(), "test.db")
			s, err := NewSQLiteStore(dbPath)
			if err != nil {
				t.Fatalf("creating SQLiteStore: %v", err)
			}
			t.Cleanup(func() { _ = s.Close() })
			return s
		})
	})
	t.Run("Memory", func(t *testing.T) {
		f(t, func(t *testing.T) Store {
			t.Helper()
			return NewMemoryStore()
		})
	})
}

func seedData(t *testing.T, store Store) {
	t.Helper()
	ctx := context.Background()

	// Login attempts: root/toor from two IPs, admin/admin from one IP.
	for range 5 {
		if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.1", ""); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}
	for range 3 {
		if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.2", ""); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}
	for range 2 {
		if err := store.RecordLoginAttempt(ctx, "admin", "admin", "10.0.0.1", ""); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}

	// Sessions: one active, one ended.
	id1, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
	if err != nil {
		t.Fatalf("creating session: %v", err)
	}
	if err := store.EndSession(ctx, id1, time.Now()); err != nil {
		t.Fatalf("ending session: %v", err)
	}

	if _, err := store.CreateSession(ctx, "10.0.0.2", "admin", "bash", ""); err != nil {
		t.Fatalf("creating session: %v", err)
	}
}

func TestGetDashboardStats(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()
			stats, err := store.GetDashboardStats(ctx)
			if err != nil {
				t.Fatalf("GetDashboardStats: %v", err)
			}
			if stats.TotalAttempts != 0 || stats.UniqueIPs != 0 || stats.TotalSessions != 0 || stats.ActiveSessions != 0 {
				t.Errorf("expected all zeros, got %+v", stats)
			}
		})

		t.Run("with data", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)
			ctx := context.Background()

			stats, err := store.GetDashboardStats(ctx)
			if err != nil {
				t.Fatalf("GetDashboardStats: %v", err)
			}
			// 5 + 3 + 2 = 10 total attempts
			if stats.TotalAttempts != 10 {
				t.Errorf("TotalAttempts = %d, want 10", stats.TotalAttempts)
			}
			// 2 unique IPs: 10.0.0.1 and 10.0.0.2
			if stats.UniqueIPs != 2 {
				t.Errorf("UniqueIPs = %d, want 2", stats.UniqueIPs)
			}
			if stats.TotalSessions != 2 {
				t.Errorf("TotalSessions = %d, want 2", stats.TotalSessions)
			}
			if stats.ActiveSessions != 1 {
				t.Errorf("ActiveSessions = %d, want 1", stats.ActiveSessions)
			}
		})
	})
}

func TestGetTopUsernames(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			entries, err := store.GetTopUsernames(context.Background(), 10)
			if err != nil {
				t.Fatalf("GetTopUsernames: %v", err)
			}
			if len(entries) != 0 {
				t.Errorf("expected empty, got %v", entries)
			}
		})

		t.Run("with data", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)

			entries, err := store.GetTopUsernames(context.Background(), 10)
			if err != nil {
				t.Fatalf("GetTopUsernames: %v", err)
			}
			if len(entries) != 2 {
				t.Fatalf("len = %d, want 2", len(entries))
			}
			// root: 5 + 3 = 8, admin: 2
			if entries[0].Value != "root" || entries[0].Count != 8 {
				t.Errorf("entries[0] = %+v, want root/8", entries[0])
			}
			if entries[1].Value != "admin" || entries[1].Count != 2 {
				t.Errorf("entries[1] = %+v, want admin/2", entries[1])
			}
		})

		t.Run("limit", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)

			entries, err := store.GetTopUsernames(context.Background(), 1)
			if err != nil {
				t.Fatalf("GetTopUsernames: %v", err)
			}
			if len(entries) != 1 {
				t.Fatalf("len = %d, want 1", len(entries))
			}
		})
	})
}

func TestGetTopPasswords(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		store := newStore(t)
		seedData(t, store)

		entries, err := store.GetTopPasswords(context.Background(), 10)
		if err != nil {
			t.Fatalf("GetTopPasswords: %v", err)
		}
		if len(entries) != 2 {
			t.Fatalf("len = %d, want 2", len(entries))
		}
		// toor: 8, admin: 2
		if entries[0].Value != "toor" || entries[0].Count != 8 {
			t.Errorf("entries[0] = %+v, want toor/8", entries[0])
		}
	})
}

func TestGetTopIPs(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		store := newStore(t)
		seedData(t, store)

		entries, err := store.GetTopIPs(context.Background(), 10)
		if err != nil {
			t.Fatalf("GetTopIPs: %v", err)
		}
		if len(entries) != 2 {
			t.Fatalf("len = %d, want 2", len(entries))
		}
		// 10.0.0.1: 5 + 2 = 7, 10.0.0.2: 3
		if entries[0].Value != "10.0.0.1" || entries[0].Count != 7 {
			t.Errorf("entries[0] = %+v, want 10.0.0.1/7", entries[0])
		}
	})
}

func TestGetSession(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("not found", func(t *testing.T) {
			store := newStore(t)
			s, err := store.GetSession(context.Background(), "nonexistent")
			if err != nil {
				t.Fatalf("GetSession: %v", err)
			}
			if s != nil {
				t.Errorf("expected nil, got %+v", s)
			}
		})

		t.Run("found", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()
			id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			s, err := store.GetSession(ctx, id)
			if err != nil {
				t.Fatalf("GetSession: %v", err)
			}
			if s == nil {
				t.Fatal("expected session, got nil")
			}
			if s.ID != id || s.IP != "10.0.0.1" || s.Username != "root" || s.ShellName != "bash" {
				t.Errorf("unexpected session: %+v", s)
			}
		})
	})
}

func TestGetSessionLogs(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		store := newStore(t)
		ctx := context.Background()

		id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
		if err != nil {
			t.Fatalf("CreateSession: %v", err)
		}
		if err := store.AppendSessionLog(ctx, id, "ls", "file1\nfile2"); err != nil {
			t.Fatalf("AppendSessionLog: %v", err)
		}
		if err := store.AppendSessionLog(ctx, id, "pwd", "/home/root"); err != nil {
			t.Fatalf("AppendSessionLog: %v", err)
		}

		logs, err := store.GetSessionLogs(ctx, id)
		if err != nil {
			t.Fatalf("GetSessionLogs: %v", err)
		}
		if len(logs) != 2 {
			t.Fatalf("len = %d, want 2", len(logs))
		}
		if logs[0].Input != "ls" {
			t.Errorf("logs[0].Input = %q, want %q", logs[0].Input, "ls")
		}
		if logs[1].Input != "pwd" {
			t.Errorf("logs[1].Input = %q, want %q", logs[1].Input, "pwd")
		}
	})
}

func TestSessionEvents(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			events, err := store.GetSessionEvents(context.Background(), "nonexistent")
			if err != nil {
				t.Fatalf("GetSessionEvents: %v", err)
			}
			if len(events) != 0 {
				t.Errorf("expected empty, got %d", len(events))
			}
		})

		t.Run("append and retrieve", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			now := time.Now().UTC()
			events := []SessionEvent{
				{SessionID: id, Timestamp: now, Direction: 0, Data: []byte("ls\n")},
				{SessionID: id, Timestamp: now.Add(100 * time.Millisecond), Direction: 1, Data: []byte("file1\nfile2\n")},
				{SessionID: id, Timestamp: now.Add(200 * time.Millisecond), Direction: 0, Data: []byte("pwd\n")},
			}
			if err := store.AppendSessionEvents(ctx, events); err != nil {
				t.Fatalf("AppendSessionEvents: %v", err)
			}

			got, err := store.GetSessionEvents(ctx, id)
			if err != nil {
				t.Fatalf("GetSessionEvents: %v", err)
			}
			if len(got) != 3 {
				t.Fatalf("len = %d, want 3", len(got))
			}
			if got[0].Direction != 0 || string(got[0].Data) != "ls\n" {
				t.Errorf("got[0] = %+v", got[0])
			}
			if got[1].Direction != 1 || string(got[1].Data) != "file1\nfile2\n" {
				t.Errorf("got[1] = %+v", got[1])
			}
		})

		t.Run("append empty", func(t *testing.T) {
			store := newStore(t)
			if err := store.AppendSessionEvents(context.Background(), nil); err != nil {
				t.Fatalf("AppendSessionEvents(nil): %v", err)
			}
		})
	})
}

func TestCloseActiveSessions(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("no active sessions", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			n, err := store.CloseActiveSessions(ctx, time.Now())
			if err != nil {
				t.Fatalf("CloseActiveSessions: %v", err)
			}
			if n != 0 {
				t.Errorf("closed %d, want 0", n)
			}
		})

		t.Run("closes only active sessions", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			// Create 3 sessions: end one, leave two active.
			id1, _ := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			store.CreateSession(ctx, "10.0.0.2", "admin", "bash", "")
			store.CreateSession(ctx, "10.0.0.3", "test", "bash", "")
			store.EndSession(ctx, id1, time.Now())

			n, err := store.CloseActiveSessions(ctx, time.Now())
			if err != nil {
				t.Fatalf("CloseActiveSessions: %v", err)
			}
			if n != 2 {
				t.Errorf("closed %d, want 2", n)
			}

			// Verify no active sessions remain.
			active, err := store.GetRecentSessions(ctx, 10, true)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(active) != 0 {
				t.Errorf("active sessions = %d, want 0", len(active))
			}
		})
	})
}

func TestSetExecCommand(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("set and retrieve", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			// Initially nil.
			s, err := store.GetSession(ctx, id)
			if err != nil {
				t.Fatalf("GetSession: %v", err)
			}
			if s.ExecCommand != nil {
				t.Errorf("expected nil ExecCommand, got %q", *s.ExecCommand)
			}

			// Set exec command.
			if err := store.SetExecCommand(ctx, id, "uname -a"); err != nil {
				t.Fatalf("SetExecCommand: %v", err)
			}

			s, err = store.GetSession(ctx, id)
			if err != nil {
				t.Fatalf("GetSession: %v", err)
			}
			if s.ExecCommand == nil {
				t.Fatal("expected non-nil ExecCommand")
			}
			if *s.ExecCommand != "uname -a" {
				t.Errorf("ExecCommand = %q, want %q", *s.ExecCommand, "uname -a")
			}
		})

		t.Run("appears in recent sessions", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			if err := store.SetExecCommand(ctx, id, "id"); err != nil {
				t.Fatalf("SetExecCommand: %v", err)
			}

			sessions, err := store.GetRecentSessions(ctx, 10, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			if sessions[0].ExecCommand == nil || *sessions[0].ExecCommand != "id" {
				t.Errorf("ExecCommand = %v, want \"id\"", sessions[0].ExecCommand)
			}
		})
	})
}

func seedChartData(t *testing.T, store Store) {
	t.Helper()
	ctx := context.Background()

	// Record attempts with country data from different IPs.
	for range 5 {
		if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.1", "CN"); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}
	for range 3 {
		if err := store.RecordLoginAttempt(ctx, "admin", "admin", "10.0.0.2", "RU"); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}
	for range 2 {
		if err := store.RecordLoginAttempt(ctx, "root", "123456", "10.0.0.3", "CN"); err != nil {
			t.Fatalf("seeding attempt: %v", err)
		}
	}
}

func TestGetAttemptsOverTime(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			points, err := store.GetAttemptsOverTime(context.Background(), 30, nil, nil)
			if err != nil {
				t.Fatalf("GetAttemptsOverTime: %v", err)
			}
			if len(points) != 0 {
				t.Errorf("expected empty, got %v", points)
			}
		})

		t.Run("with data", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			points, err := store.GetAttemptsOverTime(context.Background(), 30, nil, nil)
			if err != nil {
				t.Fatalf("GetAttemptsOverTime: %v", err)
			}
			// All data was inserted today, so should be one point.
			if len(points) != 1 {
				t.Fatalf("len = %d, want 1", len(points))
			}
			// 5 + 3 + 2 = 10 total.
			if points[0].Count != 10 {
				t.Errorf("count = %d, want 10", points[0].Count)
			}
		})
	})
}

func TestGetHourlyPattern(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			counts, err := store.GetHourlyPattern(context.Background(), nil, nil)
			if err != nil {
				t.Fatalf("GetHourlyPattern: %v", err)
			}
			if len(counts) != 0 {
				t.Errorf("expected empty, got %v", counts)
			}
		})

		t.Run("with data", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			counts, err := store.GetHourlyPattern(context.Background(), nil, nil)
			if err != nil {
				t.Fatalf("GetHourlyPattern: %v", err)
			}
			// All data was inserted at the same hour.
			if len(counts) != 1 {
				t.Fatalf("len = %d, want 1", len(counts))
			}
			if counts[0].Count != 10 {
				t.Errorf("count = %d, want 10", counts[0].Count)
			}
		})
	})
}

func TestGetCountryStats(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			counts, err := store.GetCountryStats(context.Background())
			if err != nil {
				t.Fatalf("GetCountryStats: %v", err)
			}
			if len(counts) != 0 {
				t.Errorf("expected empty, got %v", counts)
			}
		})

		t.Run("with data", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			counts, err := store.GetCountryStats(context.Background())
			if err != nil {
				t.Fatalf("GetCountryStats: %v", err)
			}
			if len(counts) != 2 {
				t.Fatalf("len = %d, want 2", len(counts))
			}
			// CN: 5 + 2 = 7, RU: 3 - ordered by count DESC.
			if counts[0].Country != "CN" || counts[0].Count != 7 {
				t.Errorf("counts[0] = %+v, want CN/7", counts[0])
			}
			if counts[1].Country != "RU" || counts[1].Count != 3 {
				t.Errorf("counts[1] = %+v, want RU/3", counts[1])
			}
		})

		t.Run("excludes empty country", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()
			if err := store.RecordLoginAttempt(ctx, "test", "test", "10.0.0.1", ""); err != nil {
				t.Fatalf("seeding: %v", err)
			}
			if err := store.RecordLoginAttempt(ctx, "test", "test", "10.0.0.2", "US"); err != nil {
				t.Fatalf("seeding: %v", err)
			}

			counts, err := store.GetCountryStats(ctx)
			if err != nil {
				t.Fatalf("GetCountryStats: %v", err)
			}
			if len(counts) != 1 {
				t.Fatalf("len = %d, want 1", len(counts))
			}
			if counts[0].Country != "US" {
				t.Errorf("country = %q, want US", counts[0].Country)
			}
		})
	})
}

func TestGetFilteredDashboardStats(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("no filter", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			stats, err := store.GetFilteredDashboardStats(context.Background(), DashboardFilter{})
			if err != nil {
				t.Fatalf("GetFilteredDashboardStats: %v", err)
			}
			if stats.TotalAttempts != 10 {
				t.Errorf("TotalAttempts = %d, want 10", stats.TotalAttempts)
			}
		})

		t.Run("filter by country", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			stats, err := store.GetFilteredDashboardStats(context.Background(), DashboardFilter{Country: "CN"})
			if err != nil {
				t.Fatalf("GetFilteredDashboardStats: %v", err)
			}
			// CN: 5 + 2 = 7
			if stats.TotalAttempts != 7 {
				t.Errorf("TotalAttempts = %d, want 7", stats.TotalAttempts)
			}
		})

		t.Run("filter by IP", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			stats, err := store.GetFilteredDashboardStats(context.Background(), DashboardFilter{IP: "10.0.0.1"})
			if err != nil {
				t.Fatalf("GetFilteredDashboardStats: %v", err)
			}
			if stats.TotalAttempts != 5 {
				t.Errorf("TotalAttempts = %d, want 5", stats.TotalAttempts)
			}
		})

		t.Run("filter by username", func(t *testing.T) {
			store := newStore(t)
			seedChartData(t, store)

			stats, err := store.GetFilteredDashboardStats(context.Background(), DashboardFilter{Username: "admin"})
			if err != nil {
				t.Fatalf("GetFilteredDashboardStats: %v", err)
			}
			if stats.TotalAttempts != 3 {
				t.Errorf("TotalAttempts = %d, want 3", stats.TotalAttempts)
			}
		})
	})
}

func TestGetFilteredTopUsernames(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		store := newStore(t)
		seedChartData(t, store)

		// Filter by country CN should only show root.
		entries, err := store.GetFilteredTopUsernames(context.Background(), 10, DashboardFilter{Country: "CN"})
		if err != nil {
			t.Fatalf("GetFilteredTopUsernames: %v", err)
		}
		if len(entries) != 1 {
			t.Fatalf("len = %d, want 1", len(entries))
		}
		if entries[0].Value != "root" || entries[0].Count != 7 {
			t.Errorf("entries[0] = %+v, want root/7", entries[0])
		}
	})
}

func TestGetRecentSessions(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("empty", func(t *testing.T) {
			store := newStore(t)
			sessions, err := store.GetRecentSessions(context.Background(), 10, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 0 {
				t.Errorf("expected empty, got %d", len(sessions))
			}
		})

		t.Run("all sessions", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)

			sessions, err := store.GetRecentSessions(context.Background(), 10, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 2 {
				t.Fatalf("len = %d, want 2", len(sessions))
			}
		})

		t.Run("active only", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)

			sessions, err := store.GetRecentSessions(context.Background(), 10, true)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			if sessions[0].DisconnectedAt != nil {
				t.Error("active session should have nil DisconnectedAt")
			}
		})

		t.Run("limit", func(t *testing.T) {
			store := newStore(t)
			seedData(t, store)

			sessions, err := store.GetRecentSessions(context.Background(), 1, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
		})
	})
}

func TestInputBytes(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("counts only input direction", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			id, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			now := time.Now().UTC()
			events := []SessionEvent{
				{SessionID: id, Timestamp: now, Direction: 0, Data: []byte("ls\n")},                           // 3 bytes input
				{SessionID: id, Timestamp: now.Add(100 * time.Millisecond), Direction: 1, Data: []byte("file1\nfile2\n")}, // 11 bytes output
				{SessionID: id, Timestamp: now.Add(200 * time.Millisecond), Direction: 0, Data: []byte("pwd\n")},          // 4 bytes input
			}
			if err := store.AppendSessionEvents(ctx, events); err != nil {
				t.Fatalf("AppendSessionEvents: %v", err)
			}

			sessions, err := store.GetRecentSessions(ctx, 10, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			// Only direction=0 data: "ls\n" (3) + "pwd\n" (4) = 7
			if sessions[0].InputBytes != 7 {
				t.Errorf("InputBytes = %d, want 7", sessions[0].InputBytes)
			}
		})

		t.Run("zero when no events", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			_, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			sessions, err := store.GetRecentSessions(ctx, 10, false)
			if err != nil {
				t.Fatalf("GetRecentSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			if sessions[0].InputBytes != 0 {
				t.Errorf("InputBytes = %d, want 0", sessions[0].InputBytes)
			}
		})
	})
}

func TestGetFilteredSessions(t *testing.T) {
	testStores(t, func(t *testing.T, newStore storeFactory) {
		t.Run("filter by human score", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			// Create two sessions, one with human score > 0.
			id1, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "CN")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			if err := store.UpdateHumanScore(ctx, id1, 0.75); err != nil {
				t.Fatalf("UpdateHumanScore: %v", err)
			}

			_, err = store.CreateSession(ctx, "10.0.0.2", "admin", "bash", "US")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			sessions, err := store.GetFilteredSessions(ctx, 50, false, DashboardFilter{HumanScoreAboveZero: true})
			if err != nil {
				t.Fatalf("GetFilteredSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			if sessions[0].ID != id1 {
				t.Errorf("expected session %s, got %s", id1, sessions[0].ID)
			}
		})

		t.Run("sort by input bytes", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			// Session with more input (created first).
			id1, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			now := time.Now().UTC()
			if err := store.AppendSessionEvents(ctx, []SessionEvent{
				{SessionID: id1, Timestamp: now, Direction: 0, Data: []byte("ls -la /tmp\n")},
				{SessionID: id1, Timestamp: now.Add(time.Millisecond), Direction: 0, Data: []byte("cat /etc/passwd\n")},
			}); err != nil {
				t.Fatalf("AppendSessionEvents: %v", err)
			}

			// Session with less input (created after id1, so would be first by connected_at).
			// Sleep >1s to ensure different RFC3339 timestamps in SQLite.
			time.Sleep(1100 * time.Millisecond)
			id2, err := store.CreateSession(ctx, "10.0.0.2", "admin", "bash", "")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			if err := store.AppendSessionEvents(ctx, []SessionEvent{
				{SessionID: id2, Timestamp: now.Add(2 * time.Second), Direction: 0, Data: []byte("x\n")},
			}); err != nil {
				t.Fatalf("AppendSessionEvents: %v", err)
			}

			// Default sort (connected_at DESC) should show id2 first.
			sessions, err := store.GetFilteredSessions(ctx, 50, false, DashboardFilter{})
			if err != nil {
				t.Fatalf("GetFilteredSessions: %v", err)
			}
			if len(sessions) != 2 {
				t.Fatalf("len = %d, want 2", len(sessions))
			}
			if sessions[0].ID != id2 {
				t.Errorf("default sort: expected %s first, got %s", id2, sessions[0].ID)
			}

			// Sort by input_bytes should show id1 first (more input).
			sessions, err = store.GetFilteredSessions(ctx, 50, false, DashboardFilter{SortBy: "input_bytes"})
			if err != nil {
				t.Fatalf("GetFilteredSessions: %v", err)
			}
			if len(sessions) != 2 {
				t.Fatalf("len = %d, want 2", len(sessions))
			}
			if sessions[0].ID != id1 {
				t.Errorf("input_bytes sort: expected %s first, got %s", id1, sessions[0].ID)
			}
		})

		t.Run("combined filters", func(t *testing.T) {
			store := newStore(t)
			ctx := context.Background()

			id1, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", "CN")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			if err := store.UpdateHumanScore(ctx, id1, 0.5); err != nil {
				t.Fatalf("UpdateHumanScore: %v", err)
			}

			// Different country, also has score.
			id2, err := store.CreateSession(ctx, "10.0.0.2", "admin", "bash", "US")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}
			if err := store.UpdateHumanScore(ctx, id2, 0.8); err != nil {
				t.Fatalf("UpdateHumanScore: %v", err)
			}

			// Same country CN but no score.
			_, err = store.CreateSession(ctx, "10.0.0.3", "test", "bash", "CN")
			if err != nil {
				t.Fatalf("CreateSession: %v", err)
			}

			// Filter: CN + human score > 0 -> only id1.
			sessions, err := store.GetFilteredSessions(ctx, 50, false, DashboardFilter{
				Country:             "CN",
				HumanScoreAboveZero: true,
			})
			if err != nil {
				t.Fatalf("GetFilteredSessions: %v", err)
			}
			if len(sessions) != 1 {
				t.Fatalf("len = %d, want 1", len(sessions))
			}
			if sessions[0].ID != id1 {
				t.Errorf("expected session %s, got %s", id1, sessions[0].ID)
			}
		})
	})
}