yoloyolo/aggregate/window_test.go

package aggregate

import (
	"testing"
	"time"

	"github.com/netsynth/netsynth/classify"
)

// sendPackets sends a slice of ClassifiedPackets to the events channel.
func sendPackets(events chan<- classify.ClassifiedPacket, packets []classify.ClassifiedPacket) {
	for _, p := range packets {
		events <- p
	}
}

func TestAggregateEmitsSnapshot(t *testing.T) {
	done := make(chan struct{})
	events := make(chan classify.ClassifiedPacket, 20)
	defer close(done)

	// Use short window to trigger ticks quickly.
	out := Aggregate(done, events, 20, nil)

	// Send 5 ICMP packets.
	for i := 0; i < 5; i++ {
		events <- classify.ClassifiedPacket{Class: classify.ClassICMP}
	}

	// Wait for at least one snapshot from a tick.
	select {
	case snap := <-out:
		if snap.TotalPackets < 5 {
			// Packets may span windows; total across first snapshot should be <= 5.
			// Just verify snapshot is emitted and counts are non-negative.
		}
		_ = snap
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: no snapshot emitted within 500ms")
	}
}

func TestAggregateMultipleClasses(t *testing.T) {
	done := make(chan struct{})
	events := make(chan classify.ClassifiedPacket, 20)

	out := Aggregate(done, events, 20, nil)

	events <- classify.ClassifiedPacket{Class: classify.ClassICMP}
	events <- classify.ClassifiedPacket{Class: classify.ClassICMP}
	events <- classify.ClassifiedPacket{Class: classify.ClassDNS}
	events <- classify.ClassifiedPacket{Class: classify.ClassHTTPS}
	events <- classify.ClassifiedPacket{Class: classify.ClassHTTPS}
	events <- classify.ClassifiedPacket{Class: classify.ClassHTTPS}

	// Close done to get the final flush.
	// But first drain any tick-based snapshots.
	time.Sleep(10 * time.Millisecond)
	close(done)

	// Accumulate all snapshots to verify totals.
	totals := make(map[classify.TrafficClass]int64)
	for snap := range out {
		for class, count := range snap.Counts {
			totals[class] += count
		}
	}

	if totals[classify.ClassICMP] != 2 {
		t.Errorf("expected ICMP=2, got %d", totals[classify.ClassICMP])
	}
	if totals[classify.ClassDNS] != 1 {
		t.Errorf("expected DNS=1, got %d", totals[classify.ClassDNS])
	}
	if totals[classify.ClassHTTPS] != 3 {
		t.Errorf("expected HTTPS=3, got %d", totals[classify.ClassHTTPS])
	}
}

func TestAggregateDoneFlushesPartial(t *testing.T) {
	done := make(chan struct{})
	// Use very large window so no tick fires.
	events := make(chan classify.ClassifiedPacket, 20)

	out := Aggregate(done, events, 100_000, nil)

	events <- classify.ClassifiedPacket{Class: classify.ClassSSH}
	events <- classify.ClassifiedPacket{Class: classify.ClassSSH}

	// Give goroutine time to process the two events.
	time.Sleep(20 * time.Millisecond)

	// Close done to trigger flush.
	close(done)

	var snap classify.WindowSnapshot
	select {
	case snap = <-out:
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: no snapshot on done-flush")
	}

	if snap.Counts[classify.ClassSSH] != 2 {
		t.Errorf("expected SSH=2 in flush, got %d", snap.Counts[classify.ClassSSH])
	}
	if snap.TotalPackets != 2 {
		t.Errorf("expected TotalPackets=2, got %d", snap.TotalPackets)
	}
}

func TestAggregateEmptyWindow(t *testing.T) {
	done := make(chan struct{})
	events := make(chan classify.ClassifiedPacket, 20)

	out := Aggregate(done, events, 20, nil)

	// Don't send any packets, just wait for a tick.
	select {
	case snap := <-out:
		if snap.TotalPackets != 0 {
			t.Errorf("expected TotalPackets=0 for empty window, got %d", snap.TotalPackets)
		}
		if len(snap.Counts) != 0 {
			t.Errorf("expected empty Counts for empty window, got %v", snap.Counts)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: no snapshot for empty window")
	}

	close(done)
}

func TestAggregateWindowIndex(t *testing.T) {
	done := make(chan struct{})
	events := make(chan classify.ClassifiedPacket, 5)

	out := Aggregate(done, events, 20, nil)

	// Collect 3 snapshots.
	var indices []int
	timeout := time.After(2 * time.Second)
	for len(indices) < 3 {
		select {
		case snap := <-out:
			indices = append(indices, snap.WindowIndex)
		case <-timeout:
			t.Fatalf("timeout: only collected %d snapshots", len(indices))
		}
	}
	close(done)

	// WindowIndex should increment: 0, 1, 2.
	for i, idx := range indices {
		if idx != i {
			t.Errorf("expected WindowIndex[%d]=%d, got %d", i, i, idx)
		}
	}
}