From 38d75b10b07c375b6a865c6e2e810cc6f9d7d27a Mon Sep 17 00:00:00 2001
From: ZGaetano <zgaetano@wilddragon.net>
Date: Sun, 10 May 2026 13:20:59 -0400
Subject: [PATCH] test(webrtc): add STAP-A IDR detection tests (issue #18)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three new test functions covering the STAP-A (NAL type 24) packetisation
mode added to isH264IDRStart:
- LeadingIDR: STAP-A where first NAL is type 5 → true
- LeadingNonIDR: STAP-A where first NAL is SPS (type 7) → false
- Truncated: STAP-A with < 4 bytes → false, no panic
---
 core/webrtc/keyframecache_test.go | 75 +++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/core/webrtc/keyframecache_test.go b/core/webrtc/keyframecache_test.go
index e11c92f..b72a3c2 100644
--- a/core/webrtc/keyframecache_test.go
+++ b/core/webrtc/keyframecache_test.go
@@ -82,6 +82,55 @@ func TestIsH264IDRStart_FUA_TruncatedPayload(t *testing.T) {
 	}
 }
 
+func TestIsH264IDRStart_STAPA_LeadingIDR(t *testing.T) {
+	// STAP-A (type 24): byte 0 = NAL header (0x78 = type 24),
+	// bytes 1-2 = first NAL size (big-endian), byte 3 = first NAL header.
+	// First NAL type = 5 (IDR) → should be detected.
+	p := makePacket([]byte{
+		0x78,       // STAP-A header: NRI=3, type=24
+		0x00, 0x03, // first NAL size = 3 bytes
+		0x65, 0x88, 0x84, // first NAL: type 5 (IDR slice)
+		0x00, 0x02, // second NAL size = 2 bytes (SPS, doesn't matter)
+		0x67, 0x42, // second NAL: SPS
+	})
+	if !isH264IDRStart(p) {
+		t.Error("STAP-A with leading IDR NAL (type 5) should be detected as IDR start")
+	}
+}
+
+func TestIsH264IDRStart_STAPA_LeadingNonIDR(t *testing.T) {
+	// STAP-A where the first NAL is SPS (type 7), not IDR.
+	// Common pattern: encoders bundle SPS+PPS+IDR in separate STAP-A,
+	// then IDR in a single NAL or FU-A. This STAP-A should not trigger reset.
+	p := makePacket([]byte{
+		0x78,       // STAP-A header
+		0x00, 0x03, // first NAL size = 3
+		0x67, 0x42, 0x00, // first NAL: SPS (type 7)
+	})
+	if isH264IDRStart(p) {
+		t.Error("STAP-A with leading SPS (type 7) should not be IDR start")
+	}
+}
+
+func TestIsH264IDRStart_STAPA_Truncated(t *testing.T) {
+	// STAP-A with fewer than 4 bytes — cannot safely read first NAL header.
+	tests := []struct {
+		name    string
+		payload []byte
+	}{
+		{"1 byte (header only)", []byte{0x78}},
+		{"2 bytes (header + 1 size byte)", []byte{0x78, 0x00}},
+		{"3 bytes (header + size, no NAL)", []byte{0x78, 0x00, 0x01}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if isH264IDRStart(makePacket(tt.payload)) {
+				t.Errorf("truncated STAP-A (%d bytes) should not panic or return true", len(tt.payload))
+			}
+		})
+	}
+}
+
 func TestIsH264IDRStart_Opus(t *testing.T) {
 	// Opus RTP payload starts with a TOC byte — definitely not H.264
 	p := makePacket([]byte{0xf8, 0xff, 0xfe})
@@ -157,6 +206,32 @@ func TestKeyFrameCache_SecondIDRResetsAgain(t *testing.T) {
 	}
 }
 
+func TestKeyFrameCache_STAPA_IDR_ResetsCache(t *testing.T) {
+	// Verify that a STAP-A with a leading IDR NAL correctly resets the burst,
+	// just like a single-NAL IDR packet does.
+	c := newKeyFrameCache()
+
+	// Pre-load some P-frames.
+	for i := 0; i < 5; i++ {
+		c.push(makePacket([]byte{0x41, byte(i)}))
+	}
+
+	stapA := makePacket([]byte{
+		0x78,       // STAP-A header
+		0x00, 0x03, // first NAL size = 3
+		0x65, 0x88, 0x84, // first NAL: IDR (type 5)
+	})
+	c.push(stapA)
+
+	snap := c.snapshot()
+	if len(snap) != 1 {
+		t.Errorf("STAP-A IDR should reset burst to 1 packet, got %d", len(snap))
+	}
+	if snap[0] != stapA {
+		t.Error("snapshot should contain the STAP-A IDR packet")
+	}
+}
+
 func TestKeyFrameCache_MaxPacketsCap(t *testing.T) {
 	c := newKeyFrameCache()
 	c.maxPackets = 5