upstream-pion-webrtc/pkg/media/h265reader/h265reader_test.go

// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
// SPDX-License-Identifier: MIT

package h265reader

import (
	"bytes"
	"io"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestH265Reader_NextNAL(t *testing.T) {
	// Test with invalid data
	reader, err := NewReader(bytes.NewReader([]byte{0xFF, 0xFF, 0xFF, 0xFF}))
	assert.NoError(t, err)

	_, err = reader.NextNAL()
	assert.Equal(t, errDataIsNotH265Stream.Error(), err.Error())

	// Test with valid H265 prefix but no NAL data
	reader, err = NewReader(bytes.NewReader([]byte{0, 0, 1}))
	assert.NoError(t, err)

	_, err = reader.NextNAL()
	assert.Equal(t, io.EOF, err)

	// Test with valid H265 NAL unit (VPS example)
	nalData := []byte{
		0x0, 0x0, 0x0, 0x1, 0x40, 0x01, 0x0C, 0x01, 0xFF, 0xFF, 0x01, 0x60, 0x00,
		0x00, 0x03, 0x00, 0x90, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x78, 0xAC, 0x09,
	}
	reader, err = NewReader(bytes.NewReader(nalData))
	assert.NoError(t, err)

	nal, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal)

	assert.Equal(t, NalUnitTypeVps, nal.NalUnitType)
	assert.False(t, nal.ForbiddenZeroBit)

	// Test reading multiple NAL units
	nalData = append(nalData, []byte{
		0x0, 0x0, 0x0, 0x1, 0x42, 0x01, 0x01, 0x01, 0x60, 0x00, 0x00, 0x03,
		0x00, 0x90, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x78, 0xA0,
		0x03, 0xC0, 0x80, 0x10, 0xE5, 0x96, 0x56, 0x69, 0x24, 0xCA, 0xE0,
		0x10, 0x00, 0x00, 0x03, 0x00, 0x10, 0x00, 0x00, 0x03, 0x01, 0xE0, 0x80,
	}...)
	reader, err = NewReader(bytes.NewReader(nalData))
	assert.NoError(t, err)

	// First NAL (VPS)
	nal1, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.Equal(t, NalUnitTypeVps, nal1.NalUnitType)

	// Second NAL (SPS)
	nal2, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.Equal(t, NalUnitTypeSps, nal2.NalUnitType)

	// Test EOF
	_, err = reader.NextNAL()
	assert.Equal(t, io.EOF, err)
}

func TestH265Reader_processByte(t *testing.T) {
	reader := &H265Reader{
		nalBuffer:                   []byte{1, 2, 3, 0, 0},
		countOfConsecutiveZeroBytes: 2,
	}

	// Test finding NAL boundary
	nalFound := reader.processByte(1)
	assert.True(t, nalFound)
	assert.Equal(t, 3, len(reader.nalBuffer))

	// Test zero byte counting
	reader.countOfConsecutiveZeroBytes = 0
	nalFound = reader.processByte(0)
	assert.False(t, nalFound)
	assert.Equal(t, 1, reader.countOfConsecutiveZeroBytes)

	// Test non-zero, non-one byte
	reader.countOfConsecutiveZeroBytes = 5
	nalFound = reader.processByte(0xFF)
	assert.False(t, nalFound)
	assert.Equal(t, 0, reader.countOfConsecutiveZeroBytes)
}

func TestH265Reader_SEIPrefixSuffixSkippedByDefault(t *testing.T) {
	// Build a small Annex-B stream: VPS, PrefixSEI, SuffixSEI, SPS
	// NAL header is 2 bytes. NAL unit type is encoded in bits 1..6 of the first byte.
	stream := []byte{
		0x0, 0x0, 0x0, 0x1, 0x40, 0x01, 0xFF, // VPS (type 32)
		0x0, 0x0, 0x0, 0x1, 0x4E, 0x01, 0xFF, // PrefixSEI (type 39)
		0x0, 0x0, 0x0, 0x1, 0x50, 0x01, 0xFF, // SuffixSEI (type 40)
		0x0, 0x0, 0x0, 0x1, 0x42, 0x01, 0xFF, // SPS (type 33)
	}

	reader, err := NewReader(bytes.NewReader(stream))
	assert.NoError(t, err)

	nal1, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal1)
	assert.Equal(t, NalUnitTypeVps, nal1.NalUnitType)

	// SEI should be skipped by default (both Prefix and Suffix)
	nal2, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal2)
	assert.Equal(t, NalUnitTypeSps, nal2.NalUnitType)

	_, err = reader.NextNAL()
	assert.Equal(t, io.EOF, err)
}

func TestH265Reader_IncludeSEI(t *testing.T) {
	vps := []byte{0x40, 0x01}       // NalUnitTypeVps (32)
	prefixSEI := []byte{0x4E, 0x01} // NalUnitTypePrefixSei (39)
	suffixSEI := []byte{0x50, 0x01} // NalUnitTypeSuffixSei (40)
	sps := []byte{0x42, 0x01}       // NalUnitTypeSps (33)
	start := []byte{0x0, 0x0, 0x0, 0x1}

	stream := append([]byte{}, start...)
	stream = append(stream, vps...)
	stream = append(stream, start...)
	stream = append(stream, prefixSEI...)
	stream = append(stream, start...)
	stream = append(stream, suffixSEI...)
	stream = append(stream, start...)
	stream = append(stream, sps...)

	reader, err := NewReaderWithOptions(bytes.NewReader(stream), WithIncludeSEI(true))
	assert.NoError(t, err)

	nal1, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal1)
	assert.Equal(t, NalUnitTypeVps, nal1.NalUnitType)

	nal2, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal2)
	assert.Equal(t, NalUnitTypePrefixSei, nal2.NalUnitType)

	nal3, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal3)
	assert.Equal(t, NalUnitTypeSuffixSei, nal3.NalUnitType)

	nal4, err := reader.NextNAL()
	assert.NoError(t, err)
	assert.NotNil(t, nal4)
	assert.Equal(t, NalUnitTypeSps, nal4.NalUnitType)
}

func TestNAL_parseHeader(t *testing.T) {
	// Test VPS NAL header parsing
	data := []byte{0x40, 0x01, 0x0C, 0x01} // VPS NAL unit
	nal := newNal(data)
	nal.parseHeader()

	assert.False(t, nal.ForbiddenZeroBit)
	assert.Equal(t, NalUnitTypeVps, nal.NalUnitType)
	assert.Equal(t, uint8(0), nal.LayerID)
	assert.Equal(t, uint8(1), nal.TemporalIDPlus1)

	// Test SPS NAL header parsing
	data = []byte{0x42, 0x01, 0x01, 0x01} // SPS NAL unit
	nal = newNal(data)
	nal.parseHeader()

	assert.False(t, nal.ForbiddenZeroBit)
	assert.Equal(t, NalUnitTypeSps, nal.NalUnitType)

	// Test with insufficient data
	data = []byte{0x40} // Only one byte
	nal = newNal(data)
	nal.parseHeader() // Should not panic

	// Test forbidden bit set
	data = []byte{0x80, 0x01} // Forbidden bit set
	nal = newNal(data)
	nal.parseHeader()

	assert.True(t, nal.ForbiddenZeroBit)
}