snapshot of https://github.com/pion/webrtc.git (main, 2026-09-15; blobs > 131072B omitted)

pkg/media/h264writer/h264writer_test.go

@@ -0,0 +1,190 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package h264writer
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	"github.com/pion/rtp"
+	"github.com/stretchr/testify/assert"
+)
+
+type writerCloser struct {
+	bytes.Buffer
+}
+
+var errClose = errors.New("close error")
+
+func (w *writerCloser) Close() error {
+	return errClose
+}
+
+func TestNewWith(t *testing.T) {
+	writer := &writerCloser{}
+	h264Writer := NewWith(writer)
+	assert.NotNil(t, h264Writer.Close())
+}
+
+func TestIsKeyFrame(t *testing.T) {
+	tests := []struct {
+		name    string
+		payload []byte
+		want    bool
+	}{
+		{
+			"When given a non-keyframe; it should return false",
+			[]byte{0x27, 0x90, 0x90},
+			false,
+		},
+		{
+			"When given a SPS packetized with STAP-A; it should return true",
+			[]byte{0x38, 0x00, 0x03, 0x27, 0x90, 0x90, 0x00, 0x05, 0x28, 0x90, 0x90, 0x90, 0x90},
+			true,
+		},
+		{
+			"When given a SPS with no packetization; it should return true",
+			[]byte{0x27, 0x90, 0x90, 0x00},
+			true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := isKeyFrame(tt.payload)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func TestWriteRTP(t *testing.T) {
+	tests := []struct {
+		name        string
+		payload     []byte
+		hasKeyFrame bool
+		wantBytes   []byte
+		wantErr     error
+		reuseWriter bool
+	}{
+		{
+			"When given an empty payload; it should return nil",
+			[]byte{},
+			false,
+			[]byte{},
+			nil,
+			false,
+		},
+		{
+			"When no keyframe is defined; it should discard the packet",
+			[]byte{0x25, 0x90, 0x90},
+			false,
+			[]byte{},
+			nil,
+			false,
+		},
+		{
+			"When a valid Single NAL Unit packet is given; it should unpack it without error",
+			[]byte{0x27, 0x90, 0x90},
+			true,
+			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90},
+			nil,
+			false,
+		},
+		{
+			"When a valid STAP-A packet is given; it should unpack it without error",
+			[]byte{0x38, 0x00, 0x03, 0x27, 0x90, 0x90, 0x00, 0x05, 0x28, 0x90, 0x90, 0x90, 0x90},
+			true,
+			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90, 0x00, 0x00, 0x00, 0x01, 0x28, 0x90, 0x90, 0x90, 0x90},
+			nil,
+			false,
+		},
+		{
+			"When a valid FU-A start packet is given; it should unpack it without error",
+			[]byte{0x3C, 0x85, 0x90, 0x90, 0x90},
+			true,
+			[]byte{},
+			nil,
+			true,
+		},
+		{
+			"When a valid FU-A end packet is given; it should unpack it without error",
+			[]byte{0x3C, 0x45, 0x90, 0x90, 0x90},
+			true,
+			[]byte{0x00, 0x00, 0x00, 0x01, 0x25, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90},
+			nil,
+			false,
+		},
+	}
+
+	var reuseWriter *bytes.Buffer
+	var reuseH264Writer *H264Writer
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			writer := &bytes.Buffer{}
+			h264Writer := &H264Writer{
+				hasKeyFrame: tt.hasKeyFrame,
+				writer:      writer,
+			}
+			if reuseWriter != nil {
+				writer = reuseWriter
+			}
+			if reuseH264Writer != nil {
+				h264Writer = reuseH264Writer
+			}
+
+			assert.Equal(t, tt.wantErr, h264Writer.WriteRTP(&rtp.Packet{
+				Payload: tt.payload,
+			}))
+			assert.True(t, bytes.Equal(tt.wantBytes, writer.Bytes()))
+
+			if !tt.reuseWriter {
+				assert.Nil(t, h264Writer.Close())
+				reuseWriter = nil
+				reuseH264Writer = nil
+			} else {
+				reuseWriter = writer
+				reuseH264Writer = h264Writer
+			}
+		})
+	}
+}
+
+type writerCounter struct {
+	writeCount int
+}
+
+func (w *writerCounter) Write([]byte) (int, error) {
+	w.writeCount++
+
+	return 0, nil
+}
+
+func (w *writerCounter) Close() error {
+	return nil
+}
+
+func TestNoZeroWrite(t *testing.T) {
+	payloads := [][]byte{
+		{0x1c, 0x80, 0x01, 0x02, 0x03},
+		{0x1c, 0x00, 0x04, 0x05, 0x06},
+		{0x1c, 0x00, 0x07, 0x08, 0x09},
+		{0x1c, 0x00, 0x10, 0x11, 0x12},
+		{0x1c, 0x40, 0x13, 0x14, 0x15},
+	}
+
+	writer := &writerCounter{}
+	h264Writer := &H264Writer{
+		hasKeyFrame: true,
+		writer:      writer,
+	}
+
+	for i := range payloads {
+		assert.NoError(t, h264Writer.WriteRTP(&rtp.Packet{
+			Payload: payloads[i],
+		}))
+	}
+	assert.Equal(t, 1, writer.writeCount)
+}