broad-announce/internal/tailhub/hub_test.go

package tailhub

import (
	"sync"
	"testing"
	"time"
)

func TestHub_PublishSubscribe(t *testing.T) {
	h := NewHub()
	s := h.Subscribe(Filter{})
	defer s.Unsubscribe()

	ev := &Event{AlertID: "a1", CompanyID: "acme-001", SourceID: "prom-prod"}
	h.Publish(ev)

	select {
	case got := <-s.C:
		if got.AlertID != "a1" {
			t.Fatalf("got %q, want a1", got.AlertID)
		}
	case <-time.After(time.Second):
		t.Fatal("timeout waiting for event")
	}
}

func TestHub_CompanyFilter(t *testing.T) {
	h := NewHub()
	all := h.Subscribe(Filter{})
	acme := h.Subscribe(Filter{CompanyID: "acme-001"})
	defer all.Unsubscribe()
	defer acme.Unsubscribe()

	h.Publish(&Event{AlertID: "a1", CompanyID: "acme-001"})
	h.Publish(&Event{AlertID: "g1", CompanyID: "globex-002"})

	// all sees both
	got := map[string]bool{}
	for i := 0; i < 2; i++ {
		select {
		case ev := <-all.C:
			got[ev.AlertID] = true
		case <-time.After(time.Second):
			t.Fatal("timeout")
		}
	}
	if !got["a1"] || !got["g1"] {
		t.Fatalf("all subscriber missing events, got %v", got)
	}

	// acme sees only a1
	select {
	case ev := <-acme.C:
		if ev.AlertID != "a1" {
			t.Fatalf("acme got %q, want a1", ev.AlertID)
		}
	case <-time.After(time.Second):
		t.Fatal("timeout")
	}
	select {
	case ev := <-acme.C:
		t.Fatalf("acme got unexpected %q (filter broken)", ev.AlertID)
	case <-time.After(50 * time.Millisecond):
		// expected: no more events
	}
}

func TestHub_DropsOnSlowConsumer(t *testing.T) {
	h := NewHub()
	s := h.Subscribe(Filter{})
	defer s.Unsubscribe()

	// Publish more than the buffer can hold. Buffer is 64.
	for i := 0; i < 100; i++ {
		h.Publish(&Event{AlertID: "x"})
	}

	// Drain — anything in the buffer is fine. We just need to
	// see Drops > 0 on the subscription.
	for i := 0; i < 64; i++ {
		select {
		case <-s.C:
		case <-time.After(time.Second):
			t.Fatalf("timeout at %d", i)
		}
	}

	// Allow the hub's lock-free drop path to register
	time.Sleep(10 * time.Millisecond)
	if s.Drops.Load() == 0 {
		t.Fatal("expected Drops > 0")
	}
	if stats := h.Stats(); stats.DroppedTotal == 0 {
		t.Fatal("expected hub.DroppedTotal > 0")
	}
}

func TestHub_UnsubscribeClosesChannel(t *testing.T) {
	h := NewHub()
	s := h.Subscribe(Filter{})

	s.Unsubscribe()
	// second unsubscribe is a no-op
	s.Unsubscribe()

	// Channel should be closed
	select {
	case _, ok := <-s.C:
		if ok {
			t.Fatal("channel still open after Unsubscribe")
		}
	case <-time.After(time.Second):
		t.Fatal("channel not closed after Unsubscribe")
	}
}

func TestHub_Concurrent(t *testing.T) {
	// Race detector: many publishers and one subscriber.
	h := NewHub()
	s := h.Subscribe(Filter{})
	defer s.Unsubscribe()

	const n = 1000
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < n/10; j++ {
				h.Publish(&Event{AlertID: "x"})
			}
		}()
	}

	// Drain concurrently
	drained := 0
	done := make(chan struct{})
	go func() {
		for {
			select {
			case <-s.C:
				drained++
			case <-done:
				return
			}
		}
	}()

	wg.Wait()
	time.Sleep(50 * time.Millisecond)
	close(done)

	// We can't assert exact drained (some dropped), but it must
	// be > 0 and ≤ n.
	if drained == 0 {
		t.Fatal("drained=0 (consumer dead?)")
	}
	if stats := h.Stats(); stats.PublishedTotal != uint64(n) {
		t.Fatalf("PublishedTotal=%d, want %d", stats.PublishedTotal, n)
	}
}