av/codec/h265/lex_test.go

/*
NAME
  lex.go

DESCRIPTION
  lex.go provides a lexer for taking h265 rtp format and lexing into access units.

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>

LICENSE
  Copyright (C) 2019 the Australian Ocean Lab (AusOcean).

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
 for more details.

  You should have received a copy of the GNU General Public License
  in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

package h265

import (
	"io"
	"testing"

	"bitbucket.org/ausocean/av/protocol/rtp"
)

type rtpReader struct {
	pkts []rtp.Pkt
	idx  int
}

func (r *rtpReader) Read(p []byte) (int, error) {
	if r.idx == len(r.pkts) {
		return 0, io.EOF
	}
	b := r.pkts[r.idx].Bytes(nil)
	copy(p, r.pkts[r.idx].Bytes(nil))
	r.idx++
	return len(b), nil
}

type destination [][]byte

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
}

func TestLex(t *testing.T) {
	const rtpVer = 2

	packets := []rtp.Pkt{
		// Singla NAL
		rtp.Pkt{
			V: rtpVer,
			Payload: []byte{
				// NAL header (type=32 VPS)
				0x40, 0x00,
				// NAL data
				0x01, 0x02, 0x03, 0x04,
			},
		},

		// Fragmentation (start packet)
		rtp.Pkt{
			V: rtpVer,
			Payload: []byte{
				// NAL header (type 49)
				0x62, 0x00,
				// FU header
				0x80,
				// FU Payload
				0x01, 0x02, 0x03,
			},
		},

		// Fragmentation (middle packet)
		rtp.Pkt{
			V: rtpVer,
			Payload: []byte{
				// NAL header (type 49)
				0x62, 0x00,
				// FU header
				0x00,
				// FU Payload
				0x04, 0x05, 0x06,
			},
		},

		// Fragmentation (end packet)
		rtp.Pkt{
			V: rtpVer,
			Payload: []byte{
				// NAL header (type 49)
				0x62, 0x00,
				// FU header
				0x40,
				// FU Payload
				0x07, 0x08, 0x09,
			},
		},

		// Aggregation.
		// This is the last bit of data for the access unit, so set M true
		rtp.Pkt{
			V: rtpVer,
			M: true,
			Payload: []byte{
				// NAL header (type 49)
				0x60, 0x00,
				// NAL1 Size
				0x00, 0x04,
				// NAL1 data
				0x01, 0x02, 0x03, 0x04,
				// NAL2 Size
				0x00, 0x04,
				// NAL2 data
				0x01, 0x02, 0x03, 0x04,
			},
		},

		// Singla NAL
		rtp.Pkt{
			V: rtpVer,
			Payload: []byte{
				// NAL header (type=32 VPS)
				0x40, 0x00,
				// NAL data
				0x01, 0x02, 0x03, 0x04,
			},
		},

		// Singla NAL
		rtp.Pkt{
			V: rtpVer,
			M: true,
			Payload: []byte{
				// NAL header (type=32 VPS)
				0x40, 0x00,
				// NAL data
				0x01, 0x02, 0x03, 0x04,
			},
		},
	}

	expect := [][]byte{
		{
			// NAL 1
			// Start code
			0x00, 0x00, 0x00, 0x01,
			// NAL header (type=32 VPS)
			0x40, 0x00,
			// NAL data
			0x01, 0x02, 0x03, 0x04,

			// NAL 2
			// Start code
			0x00, 0x00, 0x00, 0x01,
			// FU Payload
			0x01, 0x02, 0x03,
			// FU Payload
			0x04, 0x05, 0x06,
			// FU Payload
			0x07, 0x08, 0x09,

			// NAL 3
			// Start code
			0x00, 0x00, 0x00, 0x01,
			// NAL data
			0x01, 0x02, 0x03, 0x04,

			// NAL 4
			// Start code
			0x00, 0x00, 0x00, 0x01,
			// NAL2 data
			0x01, 0x02, 0x03, 0x04,
		},
		{
			// NAL 1
			//Start code
			0x00, 0x00, 0x00, 0x01,
			// NAL header (type=32 VPS)
			0x40, 0x00,
			// Data
			0x01, 0x02, 0x03, 0x04,

			// NAL 2
			//Start code
			0x00, 0x00, 0x00, 0x01,
			// NAL header (type=32 VPS)
			0x40, 0x00,
			// Data
			0x01, 0x02, 0x03, 0x04,
		},
	}

	r := &rtpReader{pkts: packets}
	d := &destination{}
	l := Lexer{UsingDON: false}
	err := l.Lex(d, r, 0)
	if err != nil {
		t.Fatalf("error lexing: %v\n", err)
	}

	for i, au := range expect {
		for j, b := range au {
			if b != [][]byte(*d)[i][j] {
				t.Fatalf("did not get expected data.\nGot: %v\nWant: %v\n", d, expect)
			}
		}
	}
}