codec/h264/lex_test.go: added test for RTPLexer.Lex(...) and also fixed some build errors

codec/h264/lex.go

@@ -40,6 +40,25 @@ import (
 	"bitbucket.org/ausocean/av/protocol/rtp"
 )
 
+// NAL types.
+const (
+	// Single nal units bounds.
+	typeSingleNALULowBound  = 1
+	typeSingleNALUHighBound = 23
+
+	// Single-time aggregation packets.
+	typeSTAPA = 24
+	typeSTAPB = 25
+
+	// Multi-time aggregation packets.
+	typeMTAP16 = 26
+	typeMTAP24 = 27
+
+	// Fragmentation packets.
+	typeFUA = 28
+	typeFUB = 29
+)
+
 var noDelay = make(chan time.Time)
 
 func init() {
@@ -138,24 +157,6 @@ outer:
 	return err
 }
 
-const (
-	// Single nal units bounds.
-	typeSingleNALULowBound  = 1
-	typeSingleNALUHighBound = 23
-
-	// Single-time aggregation packets.
-	typeSTAPA = 24
-	typeSTAPB = 25
-
-	// Multi-time aggregation packets.
-	typeMTAP16 = 26
-	typeMTAP24 = 27
-
-	// Fragmentation packets.
-	typeFUA = 28
-	typeFUB = 29
-)
-
 // Buffer sizes.
 const (
 	maxAUSize  = 100000
@@ -176,7 +177,7 @@ func NewRTPLexer() *RTPLexer {
 
 // LexFromRTP extracts H264 access units from an RTP stream. This function
 // expects that each read from src will provide a single RTP packet.
-func (l *RTPLexer) LexFromRTP(dst io.Writer, src io.Reader, delay time.Duration) error {
+func (l *RTPLexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
 	buf := make([]byte, maxRTPSize)
 	for {
 		n, err := src.Read(buf)

codec/h264/lex_test.go

@@ -7,6 +7,7 @@ DESCRIPTION
 
 AUTHOR
   Dan Kortschak <dan@ausocean.org>
+	Saxon A. Nelson-Milton <saxon@ausocean.org>
 
 LICENSE
   lex_test.go is Copyright (C) 2017 the Australian Ocean Lab (AusOcean)
@@ -30,6 +31,8 @@ LICENSE
 package h264
 
 import (
+	"io"
+	"testing"
 	"time"
 )
 
@@ -220,3 +223,143 @@ func TestH264(t *testing.T) {
 	}
 }
 */
+
+// rtpReader provides the RTP stream.
+type rtpReader struct {
+	packets [][]byte
+	idx     int
+}
+
+// Read implements io.Reader.
+func (r *rtpReader) Read(p []byte) (int, error) {
+	if r.idx == len(r.packets) {
+		return 0, io.EOF
+	}
+	b := r.packets[r.idx]
+	n := copy(p, b)
+	if n < len(r.packets[r.idx]) {
+		r.packets[r.idx] = r.packets[r.idx][n:]
+	} else {
+		r.idx++
+	}
+	return n, nil
+}
+
+// destination holds the access units extracted during the lexing process.
+type destination [][]byte
+
+// Write implements io.Writer.
+func (d *destination) Write(p []byte) (int, error) {
+	t := make([]byte, len(p))
+	copy(t, p)
+	*d = append([][]byte(*d), t)
+	return len(p), nil
+}
+
+// TestLex checks that the Lexer can correctly extract H264 access units from
+// h264 RTP stream in RTP payload format.
+func TestLex(t *testing.T) {
+	const rtpVer = 2
+
+	tests := []struct {
+		packets [][]byte
+		expect  [][]byte
+	}{
+		{
+			packets: [][]byte{
+				{ // Single NAL unit.
+					0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // NAL Data.
+				},
+				{ // Fragmentation (start packet).
+					0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeFUA,                       // FU indicator.
+					0x80 | typeSingleNALULowBound, // FU header.
+					0x01, 0x02, 0x03,              // FU payload.
+				},
+				{ // Fragmentation (middle packet)
+					0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeFUA,                // NAL indicator.
+					typeSingleNALULowBound, // FU header.
+					0x04, 0x05, 0x06,       // FU payload.
+				},
+				{ // Fragmentation (end packet)
+					0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeFUA,                       // NAL indicator.
+					0x40 | typeSingleNALULowBound, // FU header.
+					0x07, 0x08, 0x09,              // FU payload
+				},
+
+				{ // Aggregation. Make last packet of access unit => marker bit true.
+					0x80, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeSTAPA,  // NAL header.
+					0x00, 0x04, // NAL 1 size.
+					0x01, 0x02, 0x03, 0x04, // NAL 1 data.
+					0x00, 0x04, // NAL 2 size.
+					0x01, 0x02, 0x03, 0x04, // NAL 2 data.
+				},
+				// Second access unit.
+				{ // Single NAL unit.
+					0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // NAL Data.
+				},
+				{ // Single NAL. Make last packet of access unit => marker bit true.
+					0x80, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // NAL data.
+				},
+			},
+			expect: [][]byte{
+				// First access unit.
+				{
+					// NAL 1
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // NAL data.
+					// NAL 2
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					typeSingleNALULowBound,
+					0x01, 0x02, 0x03, // FU payload.
+					0x04, 0x05, 0x06, // FU payload.
+					0x07, 0x08, 0x09, // FU payload.
+					// NAL 3
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					0x01, 0x02, 0x03, 0x04, // NAL data.
+					// NAL 4
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					0x01, 0x02, 0x03, 0x04, // NAL 2 data
+				},
+				// Second access unit.
+				{
+					// NAL 1
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // Data.
+					// NAL 2
+					0x00, 0x00, 0x00, 0x01, // Start code.
+					typeSingleNALULowBound, // NAL header.
+					0x01, 0x02, 0x03, 0x04, // Data.
+				},
+			},
+		},
+	}
+
+	for testNum, test := range tests {
+		r := &rtpReader{packets: test.packets}
+		d := &destination{}
+		err := NewRTPLexer().Lex(d, r, 0)
+		if err != nil {
+			t.Fatalf("error lexing: %v\n", err)
+		}
+
+		for i, accessUnit := range test.expect {
+			for j, part := range accessUnit {
+				if part != [][]byte(*d)[i][j] {
+					t.Fatalf("did not get expected data for test: %v.\nGot: %v\nWant: %v\n", testNum, d, test.expect)
+				}
+			}
+		}
+	}
+}

protocol/rtmp/rtmp_test.go

@@ -199,7 +199,7 @@ func TestFromFrame(t *testing.T) {
 	if err != nil {
 		t.Errorf("Failed to create flv encoder with error: %v", err)
 	}
-	err = h264.Lex(flvEncoder, bytes.NewReader(videoData), time.Second/time.Duration(frameRate))
+	err = h264.LexFromBytestream(flvEncoder, bytes.NewReader(videoData), time.Second/time.Duration(frameRate))
 	if err != nil {
 		t.Errorf("Lexing failed with error: %v", err)
 	}
@@ -251,7 +251,7 @@ func TestFromFile(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to create encoder: %v", err)
 	}
-	err = h264.Lex(flvEncoder, f, time.Second/time.Duration(25))
+	err = h264.LexFromBytestream(flvEncoder, f, time.Second/time.Duration(25))
 	if err != nil {
 		t.Errorf("Lexing and encoding failed with error: %v", err)
 	}

revid/revid.go

@@ -272,7 +272,7 @@ func (r *Revid) setupPipeline(mtsEnc, flvEnc func(dst io.WriteCloser, rate int)
 	switch r.config.InputCodec {
 	case H264:
 		r.config.Logger.Log(logger.Info, pkg+"using H264 lexer")
-		r.lexTo = h264.Lex
+		r.lexTo = h264.LexFromBytestream
 	case Mjpeg:
 		r.config.Logger.Log(logger.Info, pkg+"using MJPEG lexer")
 		r.lexTo = mjpeg.Lex