datarhei-dragonfork-core/core/webrtc/source_test.go

package webrtc

import (
	"testing"
	"time"
)

// TestSourceSubscribe_PreFillFromCache verifies that a subscriber joining
// after an IDR packet has been pushed immediately receives the cached burst
// before any live packets arrive.
func TestSourceSubscribe_PreFillFromCache(t *testing.T) {
	src, err := NewSource("test-prefill", 0)
	if err != nil {
		t.Fatalf("NewSource: %v", err)
	}
	defer src.Close()

	src.EnableKeyFrameCache()
	src.Start()

	// Push directly into the cache — no need to go through UDP.
	idrPkt := makePacket([]byte{0x65, 0x88, 0x84})
	src.cache.push(idrPkt)
	for i := 0; i < 3; i++ {
		src.cache.push(makePacket([]byte{0x41, byte(i)}))
	}

	// Subscribe with a buffer big enough to hold the burst.
	ch := src.Subscribe(64)

	// Channel should already contain 4 packets — no live UDP required.
	if len(ch) != 4 {
		t.Errorf("expected 4 pre-filled packets, got %d", len(ch))
	}

	// First packet must be the IDR.
	first := <-ch
	if first.Payload[0]&0x1F != 5 {
		t.Errorf("first pre-fill packet should be IDR (type 5), got type %d", first.Payload[0]&0x1F)
	}

	src.Unsubscribe(ch)
}

// TestSourceSubscribe_NoCacheByDefault verifies that without
// EnableKeyFrameCache the channel starts empty.
func TestSourceSubscribe_NoCacheByDefault(t *testing.T) {
	src, err := NewSource("test-nocache", 0)
	if err != nil {
		t.Fatalf("NewSource: %v", err)
	}
	defer src.Close()

	src.Start()
	ch := src.Subscribe(64)

	if len(ch) != 0 {
		t.Errorf("expected empty channel without cache, got %d packets", len(ch))
	}
	src.Unsubscribe(ch)
}

// TestSourceSubscribe_PreFillStopsOnFullChannel verifies that pre-fill does
// not block when bufDepth is smaller than the burst length.
func TestSourceSubscribe_PreFillStopsOnFullChannel(t *testing.T) {
	src, err := NewSource("test-smallbuf", 0)
	if err != nil {
		t.Fatalf("NewSource: %v", err)
	}
	defer src.Close()

	src.EnableKeyFrameCache()
	src.Start()

	// Push 10 packets into the cache.
	src.cache.push(makePacket([]byte{0x65, 0x88})) // IDR
	for i := 0; i < 9; i++ {
		src.cache.push(makePacket([]byte{0x41, byte(i)}))
	}

	// Subscribe with bufDepth=3 — only 3 should land.
	ch := src.Subscribe(3)

	if len(ch) != 3 {
		t.Errorf("expected exactly 3 pre-filled packets (bufDepth cap), got %d", len(ch))
	}
	src.Unsubscribe(ch)
}

// TestSourceClose_UnsubscribesAll verifies that Close closes every subscriber
// channel so goroutines ranging over them terminate cleanly.
func TestSourceClose_UnsubscribesAll(t *testing.T) {
	src, err := NewSource("test-close", 0)
	if err != nil {
		t.Fatalf("NewSource: %v", err)
	}

	src.Start()
	ch1 := src.Subscribe(8)
	ch2 := src.Subscribe(8)

	if err := src.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	done := make(chan struct{}, 2)
	go func() {
		for range ch1 {
		}
		done <- struct{}{}
	}()
	go func() {
		for range ch2 {
		}
		done <- struct{}{}
	}()

	timeout := time.After(500 * time.Millisecond)
	for i := 0; i < 2; i++ {
		select {
		case <-done:
		case <-timeout:
			t.Error("subscriber channel not closed within 500ms of src.Close()")
			return
		}
	}
}

// TestEnableKeyFrameCache_Idempotent verifies that calling EnableKeyFrameCache
// twice does not replace or reset an existing cache.
func TestEnableKeyFrameCache_Idempotent(t *testing.T) {
	src, err := NewSource("test-idempotent", 0)
	if err != nil {
		t.Fatalf("NewSource: %v", err)
	}
	defer src.Close()

	src.EnableKeyFrameCache()
	firstCache := src.cache
	src.cache.push(makePacket([]byte{0x65, 0x01}))

	src.EnableKeyFrameCache() // second call — must be a no-op

	if src.cache != firstCache {
		t.Error("EnableKeyFrameCache should not replace an existing cache")
	}
	if len(src.cache.snapshot()) != 1 {
		t.Error("second EnableKeyFrameCache call should not clear the cache contents")
	}
}