ZLMediaKit/extension/frame_test.go

package extension

import (
	"sync"
	"testing"
	"time"
)

func TestTrackTypeConversion(t *testing.T) {
	// 验证字符串转TrackType
	if GetTrackType("video") != TrackVideo {
		t.Error("video track type conversion failed")
	}
	if GetTrackType("invalid") != TrackInvalid {
		t.Error("invalid track type conversion failed")
	}

	// 验证TrackType转字符串
	if TrackVideo.String() != "video" {
		t.Error("TrackVideo string conversion failed")
	}
}

func TestCodecIDConversion(t *testing.T) {
	// 验证名称转CodecID
	if GetCodecID("H264") != CodecH264 {
		t.Error("H264 codec conversion failed")
	}
	if GetCodecID("invalid") != CodecInvalid {
		t.Error("invalid codec conversion failed")
	}

	// 验证CodecID转名称
	if CodecH264.String() != "H264" {
		t.Error("CodecH264 string conversion failed")
	}
}

func TestBaseFrame(t *testing.T) {
	frame := &BaseFrame{
		Dts:         1000,
		Pts:         1005,
		VPrefixSize: 4,
		Codec:       CodecH264,
		Track:       TrackVideo,
	}

	if frame.DTS() != 1000 {
		t.Error("DTS mismatch")
	}
	if frame.PrefixSize() != 4 {
		t.Error("PrefixSize mismatch")
	}
}

func TestFrameCacheAble(t *testing.T) {
	// 测试原始帧不可缓存
	raw := &FrameFromBytes{
		BaseFrame: &BaseFrame{Dts: 1000},
		DataPtr:   []byte{1, 2, 3},
	}
	if raw.CacheAble() {
		t.Error("raw frame should not be cacheable")
	}

	// 测试转换后可缓存
	cached := NewFrameCacheAble(raw, false)
	if !cached.CacheAble() {
		t.Error("cached frame should be cacheable")
	}
	if cached.Size() != 3 {
		t.Error("data size mismatch")
	}
}

func TestFrameMerger_H264Prefix(t *testing.T) {
	merger := NewFrameMerger(MergeH264Prefix)
	var output []byte

	// 创建测试帧（无前缀）
	frame1 := &FrameFromBytes{BaseFrame: &BaseFrame{Dts: 1000, VPrefixSize: 0}, DataPtr: []byte{0x67}}
	frame2 := &FrameFromBytes{BaseFrame: &BaseFrame{Dts: 1000, VPrefixSize: 0}, DataPtr: []byte{0x68}}

	// 合并相同时间戳的帧
	merger.InputFrame(frame1, func(dts, pts uint64, buf []byte, key bool) {
		output = append(output, buf...)
	})
	merger.InputFrame(frame2, func(dts, pts uint64, buf []byte, key bool) {
		output = append(output, buf...)
	})
	merger.Flush(func(dts, pts uint64, buf []byte, key bool) {
		output = append(output, buf...)
	})

	// 验证H264前缀添加
	expected := []byte{0, 0, 0, 1, 0x67, 0, 0, 0, 1, 0x68}
	if string(output) != string(expected) {
		t.Error("H264 prefix merge failed")
	}
}

func TestFrameMerger_MP4NalSize(t *testing.T) {
	merger := NewFrameMerger(MergeMP4NalSize)
	var output []byte

	// 创建测试帧（带4字节前缀）
	frame := &FrameFromBytes{
		BaseFrame: &BaseFrame{VPrefixSize: 4, Dts: 1000},
		DataPtr:   []byte{0, 0, 0, 1, 0x67},
	}

	merger.InputFrame(frame, func(dts, pts uint64, buf []byte, key bool) {
		output = buf
	})
	merger.Flush(func(dts, pts uint64, buf []byte, key bool) {
		output = buf
	})

	// 验证MP4 NALU大小（4字节长度 + 原始数据）
	if len(output) != 5 {
		t.Error("MP4 nal size length mismatch")
	}
	if output[0] != 0 || output[1] != 0 || output[2] != 0 || output[3] != 1 {
		t.Error("MP4 nal size header mismatch")
	}
}

func TestFrameDispatcher(t *testing.T) {
	dispatcher := NewFrameDispatcher()
	statsChan := make(chan struct{})

	// 添加测试代理
	dispatcher.AddDelegate(&FrameWriter{
		OnFrame: func(frame Frame) bool {
			if frame.KeyFrame() && frame.TrackType() == TrackVideo {
				close(statsChan)
			}
			return true
		},
	})

	// 输入视频关键帧
	dispatcher.InputFrame(&FrameFromBytes{
		BaseFrame: &BaseFrame{
			Dts:   1000,
			Track: TrackVideo,
		},
		IsKey: true,
	})

	// 验证GOP统计
	select {
	case <-statsChan:
		frames, keys, gopSize, _ := dispatcher.GetStats()
		if frames != 1 || keys != 1 || gopSize != 1 {
			t.Error("GOP statistics mismatch")
		}
	case <-time.After(100 * time.Millisecond):
		t.Error("timeout waiting for stats")
	}
}

func TestFrameDispatcher_Concurrent(t *testing.T) {
	dispatcher := NewFrameDispatcher()
	var done sync.WaitGroup

	// 并发添加/删除代理
	for i := 0; i < 10; i++ {
		done.Add(1)
		go func(id int) {
			defer done.Done()
			writer := &FrameWriter{}
			dispatcher.AddDelegate(writer)
			dispatcher.DelDelegate(writer)
		}(i)
	}
	done.Wait()

	// 验证线程安全
	if dispatcher.Size() != 0 {
		t.Error("concurrent delegate management failed")
	}
}