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 i := 0; i < 5; i++ { if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.1"); err != nil { t.Fatalf("seeding attempt: %v", err) } } for i := 0; i < 3; i++ { if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.2"); err != nil { t.Fatalf("seeding attempt: %v", err) } } for i := 0; i < 2; i++ { 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 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)) } }) }) }