From 15af16ce972480212aada286fd45e787eca81607 Mon Sep 17 00:00:00 2001
From: ZGaetano <zgaetano@wilddragon.net>
Date: Wed, 6 May 2026 15:56:43 -0400
Subject: [PATCH] test(prometheus): add WebRTC snapshot collector unit tests

---
 prometheus/webrtc_test.go | 97 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 prometheus/webrtc_test.go

diff --git a/prometheus/webrtc_test.go b/prometheus/webrtc_test.go
new file mode 100644
index 0000000..bd5ccb8
--- /dev/null
+++ b/prometheus/webrtc_test.go
@@ -0,0 +1,97 @@
+package prometheus_test
+
+import (
+	"strings"
+	"testing"
+
+	coreprom "github.com/datarhei/core/v16/prometheus"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+)
+
+// fakeStats implements coreprom.WebRTCStatsSource for testing.
+type fakeStats struct{ s coreprom.WebRTCStats }
+
+func (f *fakeStats) Stats() coreprom.WebRTCStats { return f.s }
+
+func newReg(t *testing.T, source coreprom.WebRTCStatsSource) *prometheus.Registry {
+	t.Helper()
+	reg := prometheus.NewRegistry()
+	if err := reg.Register(coreprom.NewWebRTCCollector("test", source)); err != nil {
+		t.Fatalf("Register: %v", err)
+	}
+	return reg
+}
+
+func TestWebRTCCollector_NoStreams(t *testing.T) {
+	reg := newReg(t, &fakeStats{})
+	if err := testutil.GatherAndCompare(reg, strings.NewReader(`
+# HELP dragonfork_webrtc_active_streams Streams currently registered (processes with webrtc.enabled=true running).
+# TYPE dragonfork_webrtc_active_streams gauge
+dragonfork_webrtc_active_streams{core="test"} 0
+# HELP dragonfork_webrtc_udp_ports_in_use UDP ports currently allocated (2 per active stream).
+# TYPE dragonfork_webrtc_udp_ports_in_use gauge
+dragonfork_webrtc_udp_ports_in_use{core="test"} 0
+`),
+		"dragonfork_webrtc_active_streams",
+		"dragonfork_webrtc_udp_ports_in_use",
+	); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestWebRTCCollector_OneStreamWithPeers(t *testing.T) {
+	src := &fakeStats{s: coreprom.WebRTCStats{
+		StreamCount:   1,
+		PeersByStream: map[string]int{"live": 3},
+		UDPPortsInUse: 2,
+	}}
+	reg := newReg(t, src)
+	if err := testutil.GatherAndCompare(reg, strings.NewReader(`
+# HELP dragonfork_webrtc_active_peers Currently subscribed WHEP peers per stream.
+# TYPE dragonfork_webrtc_active_peers gauge
+dragonfork_webrtc_active_peers{core="test",stream_id="live"} 3
+# HELP dragonfork_webrtc_active_streams Streams currently registered (processes with webrtc.enabled=true running).
+# TYPE dragonfork_webrtc_active_streams gauge
+dragonfork_webrtc_active_streams{core="test"} 1
+# HELP dragonfork_webrtc_udp_ports_in_use UDP ports currently allocated (2 per active stream).
+# TYPE dragonfork_webrtc_udp_ports_in_use gauge
+dragonfork_webrtc_udp_ports_in_use{core="test"} 2
+`)); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestWebRTCCollector_MultipleStreams(t *testing.T) {
+	src := &fakeStats{s: coreprom.WebRTCStats{
+		StreamCount:   2,
+		PeersByStream: map[string]int{"live": 3, "cam": 1},
+		UDPPortsInUse: 4,
+	}}
+	reg := newReg(t, src)
+	mfs, err := reg.Gather()
+	if err != nil {
+		t.Fatalf("Gather: %v", err)
+	}
+	// Check stream count and udp ports
+	for _, mf := range mfs {
+		switch mf.GetName() {
+		case "dragonfork_webrtc_active_streams":
+			if got := mf.GetMetric()[0].GetGauge().GetValue(); got != 2 {
+				t.Errorf("active_streams: want 2, got %v", got)
+			}
+		case "dragonfork_webrtc_udp_ports_in_use":
+			if got := mf.GetMetric()[0].GetGauge().GetValue(); got != 4 {
+				t.Errorf("udp_ports_in_use: want 4, got %v", got)
+			}
+		case "dragonfork_webrtc_active_peers":
+			total := 0.0
+			for _, m := range mf.GetMetric() {
+				total += m.GetGauge().GetValue()
+			}
+			if total != 4 {
+				t.Errorf("active_peers total: want 4, got %v", total)
+			}
+		}
+	}
+}