From 8ea759e4fe10893d5fa4a0436e953649c6ca1251 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torbjo=CC=88rn=20Einarsson?= <torbjorn@einarssons.se>
Date: Tue, 5 Nov 2024 19:19:14 +0100
Subject: [PATCH] test: add more unit tests to avc package

---
 avc/annexb_test.go           | 12 ++++++++++++
 avc/avc.go                   |  8 ++++----
 avc/avc_test.go              | 22 ++++++++++++++++++++++
 avc/avcdecoderconfig_test.go | 27 +++++++++++++++++++++++++++
 avc/nalus.go                 |  6 +++---
 5 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/avc/annexb_test.go b/avc/annexb_test.go
index efa4c291..74f71883 100644
--- a/avc/annexb_test.go
+++ b/avc/annexb_test.go
@@ -3,6 +3,7 @@ package avc
 import (
 	"bytes"
 	"math/rand"
+	"os"
 	"testing"
 
 	"github.com/go-test/deep"
@@ -233,6 +234,17 @@ func TestGetFirstAVCVideoNALUFromByteStream(t *testing.T) {
 	}
 }
 
+func TestConvertByteStreamToNaluSample(t *testing.T) {
+	data, err := os.ReadFile("testdata/two-frames.264")
+	if err != nil {
+		t.Fatalf("Read: %v", err)
+	}
+	sample := ConvertByteStreamToNaluSample(data)
+	if len(sample) != len(data) {
+		t.Errorf("got %d, wanted %d", len(sample), len(data))
+	}
+}
+
 func BenchmarkByteStreamToNaluSample(b *testing.B) {
 	l := 1024 * 1024
 	data := make([]byte, l)
diff --git a/avc/avc.go b/avc/avc.go
index 3239f4a5..c982ca0a 100644
--- a/avc/avc.go
+++ b/avc/avc.go
@@ -61,11 +61,11 @@ func GetNaluType(naluHeader byte) NaluType {
 
 // FindNaluTypes - find list of NAL unit types in sample
 func FindNaluTypes(sample []byte) []NaluType {
-	naluList := make([]NaluType, 0)
 	length := len(sample)
 	if length < 4 {
-		return naluList
+		return nil
 	}
+	naluList := make([]NaluType, 0, 2)
 	var pos uint32 = 0
 	for pos < uint32(length-4) {
 		naluLength := binary.BigEndian.Uint32(sample[pos : pos+4])
@@ -79,11 +79,11 @@ func FindNaluTypes(sample []byte) []NaluType {
 
 // FindNaluTypesUpToFirstVideoNALU - find list of NAL unit types in sample
 func FindNaluTypesUpToFirstVideoNALU(sample []byte) []NaluType {
-	naluList := make([]NaluType, 0)
 	length := len(sample)
 	if length < 4 {
-		return naluList
+		return nil
 	}
+	naluList := make([]NaluType, 0)
 	var pos uint32 = 0
 	for pos < uint32(length-4) {
 		naluLength := binary.BigEndian.Uint32(sample[pos : pos+4])
diff --git a/avc/avc_test.go b/avc/avc_test.go
index ff90e4f0..277b630e 100644
--- a/avc/avc_test.go
+++ b/avc/avc_test.go
@@ -1,6 +1,7 @@
 package avc
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/go-test/deep"
@@ -29,6 +30,27 @@ func TestGetNaluTypes(t *testing.T) {
 		if diff := deep.Equal(got, tc.wanted); diff != nil {
 			t.Errorf("%s: %v", tc.name, diff)
 		}
+		for _, w := range tc.wanted {
+			if !ContainsNaluType(tc.input, w) {
+				t.Errorf("%s: wanted %v in %v", tc.name, w, tc.input)
+			}
+			if w == NALU_IDR && !IsIDRSample(tc.input) {
+				t.Errorf("%s: wanted IDR in %v", tc.name, tc.input)
+			}
+		}
+	}
+}
+
+func TestNaluTypeStrings(t *testing.T) {
+	named := make([]int, 0, 9)
+	for i := 0; i < 14; i++ {
+		nt := NaluType(i)
+		if !strings.HasPrefix(nt.String(), "Other") {
+			named = append(named, i)
+		}
+	}
+	if len(named) != 9 {
+		t.Errorf("Expected 9 named NaluTypes, got %d", len(named))
 	}
 }
 
diff --git a/avc/avcdecoderconfig_test.go b/avc/avcdecoderconfig_test.go
index 2a415a09..96f7f137 100644
--- a/avc/avcdecoderconfig_test.go
+++ b/avc/avcdecoderconfig_test.go
@@ -1,7 +1,9 @@
 package avc
 
 import (
+	"bytes"
 	"encoding/hex"
+	"os"
 	"testing"
 
 	"github.com/go-test/deep"
@@ -35,4 +37,29 @@ func TestAvcDecoderConfigRecord(t *testing.T) {
 	if diff := deep.Equal(got, wanted); diff != nil {
 		t.Error(diff)
 	}
+
+	enc := bytes.Buffer{}
+	err = got.Encode(&enc)
+	if err != nil {
+		t.Error("Error encoding AVCDecoderConfigurationRecord")
+	}
+	if !bytes.Equal(enc.Bytes(), byteData) {
+		t.Error("Error encoding AVCDecoderConfigurationRecord")
+	}
+}
+
+func TestCreateAVCDecConfRec(t *testing.T) {
+	data, err := os.ReadFile("testdata/blackframe.264")
+	if err != nil {
+		t.Error("Error reading file")
+	}
+	spss := ExtractNalusOfTypeFromByteStream(NALU_SPS, data, true)
+	ppss := ExtractNalusOfTypeFromByteStream(NALU_PPS, data, true)
+	if len(spss) != 1 || len(ppss) != 1 {
+		t.Error("Error extracting SPS/PPS")
+	}
+	_, err = CreateAVCDecConfRec(spss, ppss, true)
+	if err != nil {
+		t.Error("Error creating AVCDecoderConfigurationRecord")
+	}
 }
diff --git a/avc/nalus.go b/avc/nalus.go
index a144b822..8c53ecce 100644
--- a/avc/nalus.go
+++ b/avc/nalus.go
@@ -7,17 +7,17 @@ import (
 
 // GetNalusFromSample - get nalus by following 4 byte length fields
 func GetNalusFromSample(sample []byte) ([][]byte, error) {
-	naluList := make([][]byte, 0)
 	length := len(sample)
 	if length < 4 {
-		return naluList, fmt.Errorf("Less than 4 bytes, No NALUs")
+		return nil, fmt.Errorf("less than 4 bytes, No NALUs")
 	}
+	naluList := make([][]byte, 0, 2)
 	var pos uint32 = 0
 	for pos < uint32(length-4) {
 		naluLength := binary.BigEndian.Uint32(sample[pos : pos+4])
 		pos += 4
 		if int(pos+naluLength) > len(sample) {
-			return nil, fmt.Errorf("NAL length fields are bad. Not video?")
+			return nil, fmt.Errorf("NALU length fields are bad. Not video?")
 		}
 		naluList = append(naluList, sample[pos:pos+naluLength])
 		pos += naluLength