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 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))
			}
		})
	})
}