internal/agg/agg_test.go - master - rustwizard/pgtrace

internal/agg/agg_test.go
130 строк3 KB
Rust Wizard
feat(offcpu): aggregate wait types per query and expose in console/agent
08 авг 2026, 09:51
08 авг 2026, 09:5173d29f7
package agg

import (
	"testing"
	"time"

	"github.com/rustwizard/pgtrace/internal/bpf"
)

func TestStats_Add(t *testing.T) {
	t.Parallel()

	s := &Stats{Queries: make(map[string]int)}
	s.Add(5*time.Millisecond, "SELECT 1")
	s.Add(3*time.Millisecond, "SELECT 1")
	s.Add(10*time.Millisecond, "UPDATE t")

	if s.Count != 3 {
		t.Fatalf("Count = %d, want 3", s.Count)
	}

	if s.TotalMs != 18 {
		t.Fatalf("TotalMs = %f, want 18", s.TotalMs)
	}

	if s.MaxMs != 10 {
		t.Fatalf("MaxMs = %f, want 10", s.MaxMs)
	}

	if s.Queries["SELECT 1"] != 2 {
		t.Fatalf("Queries[SELECT 1] = %d, want 2", s.Queries["SELECT 1"])
	}
}

func TestStats_TopQueries(t *testing.T) {
	t.Parallel()

	s := &Stats{Queries: map[string]int{
		"SELECT a": 5,
		"UPDATE b": 9,
		"INSERT c": 2,
	}}

	got := s.TopQueries(2)

	want := []QueryEntry{{Query: "UPDATE b", Count: 9}, {Query: "SELECT a", Count: 5}}
	if len(got) != len(want) {
		t.Fatalf("len = %d, want %d", len(got), len(want))
	}

	for i := range want {
		if got[i] != want[i] {
			t.Errorf("TopQueries[%d] = %+v, want %+v", i, got[i], want[i])
		}
	}
}

func TestStats_TopQueries_nTooLarge(t *testing.T) {
	t.Parallel()

	s := &Stats{Queries: map[string]int{"a": 1, "b": 2}}

	got := s.TopQueries(10)
	if len(got) != 2 {
		t.Fatalf("len = %d, want 2 (clamped)", len(got))
	}
}

func TestStats_TopQueries_zero(t *testing.T) {
	t.Parallel()

	s := &Stats{Queries: map[string]int{"a": 1}}

	if got := s.TopQueries(0); got != nil {
		t.Fatalf("TopQueries(0) = %v, want nil", got)
	}
}

func TestAggregator_Snapshot(t *testing.T) {
	t.Parallel()

	start := time.Now()
	a := New(start)
	a.Add(bpf.Event{PID: 1, SyscallID: 0, DurationNs: 5e6}, "SELECT 1")

	_, _, stats, waits := a.Snapshot(start.Add(10 * time.Second))
	if len(stats) != 1 {
		t.Fatalf("snapshot has %d syscalls, want 1", len(stats))
	}

	if stats[0].Count != 1 {
		t.Fatalf("Count = %d, want 1", stats[0].Count)
	}

	if len(waits) != 0 {
		t.Fatalf("waits = %d, want 0", len(waits))
	}

	_, _, again, againWaits := a.Snapshot(start.Add(20 * time.Second))
	if again != nil {
		t.Fatalf("second snapshot = %v, want nil (window reset)", again)
	}

	if againWaits != nil {
		t.Fatalf("second snapshot waits = %v, want nil", againWaits)
	}
}

func TestAggregator_AddWait(t *testing.T) {
	t.Parallel()

	start := time.Now()
	a := New(start)
	a.AddWait("lwlock", 20*time.Millisecond, "UPDATE t SET x=1")
	a.AddWait("io", 30*time.Millisecond, "UPDATE t SET x=1")

	_, _, _, waits := a.Snapshot(start.Add(10 * time.Second))
	if len(waits) != 2 {
		t.Fatalf("waits = %d, want 2", len(waits))
	}

	lw := waits["lwlock"]
	if lw.TotalMs != 20 {
		t.Fatalf("lwlock TotalMs = %f, want 20", lw.TotalMs)
	}

	if lw.Queries["UPDATE t SET x=1"] != 20 {
		t.Fatalf("lwlock query ms = %f, want 20", lw.Queries["UPDATE t SET x=1"])
	}
}