From d5b10fdd6671ab68a9302976a2304197e93a330f Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sat, 14 Aug 2021 16:12:35 +0200
Subject: [PATCH] hls: improve tests

---
 internal/hls/muxer_test.go | 61 ++++++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 6 deletions(-)

diff --git a/internal/hls/muxer_test.go b/internal/hls/muxer_test.go
index 07e0eeaa..6b89c277 100644
--- a/internal/hls/muxer_test.go
+++ b/internal/hls/muxer_test.go
@@ -2,6 +2,7 @@ package hls
 
 import (
 	"io/ioutil"
+	"regexp"
 	"testing"
 	"time"
 
@@ -9,6 +10,12 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func checkTSPacket(t *testing.T, byts []byte, pid int, afc int) {
+	require.Equal(t, byte(0x47), byts[0])                                      // sync bit
+	require.Equal(t, uint16(pid), (uint16(byts[1])<<8|uint16(byts[2]))&0x1fff) // PID
+	require.Equal(t, uint8(afc), (byts[3]>>4)&0x03)                            // adaptation field control
+}
+
 func TestMuxer(t *testing.T) {
 	videoTrack, err := gortsplib.NewTrackH264(96, []byte{0x01, 0x02, 0x03, 0x04}, []byte{0x01, 0x02, 0x03, 0x04})
 	require.NoError(t, err)
@@ -16,7 +23,7 @@ func TestMuxer(t *testing.T) {
 	audioTrack, err := gortsplib.NewTrackAAC(97, []byte{17, 144})
 	require.NoError(t, err)
 
-	m, err := NewMuxer(3, 5*time.Second, videoTrack, audioTrack)
+	m, err := NewMuxer(3, 1*time.Second, videoTrack, audioTrack)
 	require.NoError(t, err)
 	defer m.Close()
 
@@ -29,8 +36,10 @@ func TestMuxer(t *testing.T) {
 
 	// group with IDR
 	err = m.WriteH264(2*time.Second, [][]byte{
-		{0x05},
-		{0x06},
+		{5}, // IDR
+		{9}, // AUD
+		{8}, // PPS
+		{7}, // SPS
 	})
 	require.NoError(t, err)
 
@@ -42,13 +51,53 @@ func TestMuxer(t *testing.T) {
 
 	// group without IDR
 	err = m.WriteH264(4*time.Second, [][]byte{
-		{0x06},
-		{0x07},
+		{6},
+		{7},
+	})
+	require.NoError(t, err)
+
+	time.Sleep(2 * time.Second)
+
+	// group with IDR
+	err = m.WriteH264(6*time.Second, [][]byte{
+		{5}, // IDR
 	})
 	require.NoError(t, err)
 
 	byts, err := ioutil.ReadAll(m.Playlist())
 	require.NoError(t, err)
 
-	require.Regexp(t, `^#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-ALLOW-CACHE:NO\n#EXT-X-TARGETDURATION:5\n#EXT-X-MEDIA-SEQUENCE:0\n#EXTINF:2,\n[0-9]+\.ts\n$`, string(byts))
+	re := regexp.MustCompile(`^#EXTM3U\n` +
+		`#EXT-X-VERSION:3\n` +
+		`#EXT-X-ALLOW-CACHE:NO\n` +
+		`#EXT-X-TARGETDURATION:2\n` +
+		`#EXT-X-MEDIA-SEQUENCE:0\n` +
+		`#EXTINF:2,\n` +
+		`([0-9]+\.ts)\n` +
+		`#EXTINF:0,\n` +
+		`([0-9]+\.ts)\n$`)
+	ma := re.FindStringSubmatch(string(byts))
+	require.NotEqual(t, nil, ma)
+
+	byts, err = ioutil.ReadAll(m.TSFile(ma[1]))
+	require.NoError(t, err)
+
+	checkTSPacket(t, byts, 0, 1)
+	byts = byts[188:]
+	checkTSPacket(t, byts, 4096, 1)
+	byts = byts[188:]
+
+	checkTSPacket(t, byts, 256, 3)
+	alen := int(byts[4])
+	byts = byts[4+alen+20:]
+	require.Equal(t,
+		[]byte{
+			0, 0, 0, 1, 9, 240,
+			0, 0, 0, 1, 5,
+			0, 0, 0, 1, 9,
+			0, 0, 0, 1, 8,
+			0, 0, 0, 1, 7,
+		},
+		byts[:26],
+	)
 }