oubliette/internal/metrics/metrics_test.go

package metrics

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"code.t-juice.club/torjus/oubliette/internal/storage"
)

func TestNew(t *testing.T) {
	m := New("1.2.3")

	// Gather all metrics and check expected names exist.
	families, err := m.registry.Gather()
	if err != nil {
		t.Fatalf("gather: %v", err)
	}

	want := map[string]bool{
		"oubliette_ssh_connections_total":    false,
		"oubliette_ssh_connections_active":   false,
		"oubliette_auth_attempts_total":      false,
		"oubliette_commands_executed_total":  false,
		"oubliette_human_score":             false,
		"oubliette_sessions_total":          false,
		"oubliette_sessions_active":         false,
		"oubliette_session_duration_seconds": false,
		"oubliette_build_info":              false,
	}

	for _, f := range families {
		if _, ok := want[f.GetName()]; ok {
			want[f.GetName()] = true
		}
	}

	for name, found := range want {
		if !found {
			t.Errorf("metric %q not registered", name)
		}
	}
}

func TestAuthAttemptsByCountry(t *testing.T) {
	m := New("1.0.0")
	m.AuthAttemptsByCountry.WithLabelValues("US").Inc()
	m.AuthAttemptsByCountry.WithLabelValues("DE").Inc()
	m.AuthAttemptsByCountry.WithLabelValues("US").Inc()

	families, err := m.registry.Gather()
	if err != nil {
		t.Fatalf("gather: %v", err)
	}

	var found bool
	for _, f := range families {
		if f.GetName() == "oubliette_auth_attempts_by_country_total" {
			found = true
			if len(f.GetMetric()) != 2 {
				t.Errorf("expected 2 label pairs (US, DE), got %d", len(f.GetMetric()))
			}
		}
	}
	if !found {
		t.Error("oubliette_auth_attempts_by_country_total not found after incrementing")
	}
}

func TestHandler(t *testing.T) {
	m := New("1.2.3")

	req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
	w := httptest.NewRecorder()
	m.Handler().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Errorf("status = %d, want 200", w.Code)
	}

	body, err := io.ReadAll(w.Body)
	if err != nil {
		t.Fatalf("reading body: %v", err)
	}

	if !strings.Contains(string(body), `oubliette_build_info{version="1.2.3"} 1`) {
		t.Errorf("response should contain build_info metric, got:\n%s", body)
	}
}

func TestStoreCollector(t *testing.T) {
	store := storage.NewMemoryStore()
	ctx := context.Background()

	// Seed some data.
	if err := store.RecordLoginAttempt(ctx, "root", "toor", "10.0.0.1", ""); err != nil {
		t.Fatalf("RecordLoginAttempt: %v", err)
	}
	if err := store.RecordLoginAttempt(ctx, "admin", "admin", "10.0.0.2", ""); err != nil {
		t.Fatalf("RecordLoginAttempt: %v", err)
	}
	if _, err := store.CreateSession(ctx, "10.0.0.1", "root", "bash", ""); err != nil {
		t.Fatalf("CreateSession: %v", err)
	}

	m := New("test")
	m.RegisterStoreCollector(store)

	families, err := m.registry.Gather()
	if err != nil {
		t.Fatalf("gather: %v", err)
	}

	wantMetrics := map[string]float64{
		"oubliette_storage_login_attempts_total": 2,
		"oubliette_storage_unique_ips":           2,
		"oubliette_storage_sessions_total":       1,
	}

	for _, f := range families {
		expected, ok := wantMetrics[f.GetName()]
		if !ok {
			continue
		}
		if len(f.GetMetric()) == 0 {
			t.Errorf("metric %q has no samples", f.GetName())
			continue
		}
		got := f.GetMetric()[0].GetGauge().GetValue()
		if got != expected {
			t.Errorf("metric %q = %f, want %f", f.GetName(), got, expected)
		}
		delete(wantMetrics, f.GetName())
	}

	for name := range wantMetrics {
		t.Errorf("metric %q not found in gather output", name)
	}
}