av/container/mts/mpegts_test.go

/*
NAME
  mpegts_test.go

DESCRIPTION
  mpegts_test.go contains testing for functionality found in mpegts.go.

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 mts

import (
	"bytes"
	"testing"

	"github.com/Comcast/gots/packet"
)

// TestBytes checks that Packet.Bytes() correctly produces a []byte
// representation of a Packet.
func TestBytes(t *testing.T) {
	const payloadLen, payloadChar, stuffingChar = 120, 0x11, 0xff
	const stuffingLen = PacketSize - payloadLen - 12

	tests := []struct {
		packet         Packet
		expectedHeader []byte
	}{
		{
			packet: Packet{
				PUSI: true,
				PID:  1,
				RAI:  true,
				CC:   4,
				AFC:  HasPayload | HasAdaptationField,
				PCRF: true,
				PCR:  1,
			},
			expectedHeader: []byte{
				0x47,                               // Sync byte.
				0x40,                               // TEI=0, PUSI=1, TP=0, PID=00000.
				0x01,                               // PID(Cont)=00000001.
				0x34,                               // TSC=00, AFC=11(adaptation followed by payload), CC=0100(4).
				byte(7 + stuffingLen),              // AFL=.
				0x50,                               // DI=0,RAI=1,ESPI=0,PCRF=1,OPCRF=0,SPF=0,TPDF=0, AFEF=0.
				0x00, 0x00, 0x00, 0x00, 0x80, 0x00, // PCR.
			},
		},
	}

	for testNum, test := range tests {
		// Construct payload.
		payload := make([]byte, 0, payloadLen)
		for i := 0; i < payloadLen; i++ {
			payload = append(payload, payloadChar)
		}

		// Fill the packet payload.
		test.packet.FillPayload(payload)

		// Create expected packet data and copy in expected header.
		expected := make([]byte, len(test.expectedHeader), PacketSize)
		copy(expected, test.expectedHeader)

		// Append stuffing.
		for i := 0; i < stuffingLen; i++ {
			expected = append(expected, stuffingChar)
		}

		// Append payload to expected bytes.
		expected = append(expected, payload...)

		// Compare got with expected.
		got := test.packet.Bytes(nil)
		if !bytes.Equal(got, expected) {
			t.Errorf("did not get expected result for test: %v.\n Got: %v\n Want: %v\n", testNum, got, expected)
		}
	}
}

// TestFindPid checks that FindPid can correctly extract the first instance
// of a PID from an MPEG-TS stream.
func TestFindPid(t *testing.T) {
	const targetPacketNum, numOfPackets, targetPid, stdPid = 6, 15, 1, 0

	// Prepare the stream of packets.
	var stream []byte
	for i := 0; i < numOfPackets; i++ {
		pid := uint16(stdPid)
		if i == targetPacketNum {
			pid = targetPid
		}

		p := Packet{
			PID: pid,
			AFC: hasPayload | hasAdaptationField,
		}
		p.FillPayload([]byte{byte(i)})
		stream = append(stream, p.Bytes(nil)...)
	}

	// Try to find the targetPid in the stream.
	p, i, err := FindPid(stream, targetPid)
	if err != nil {
		t.Fatalf("unexpected error finding PID: %v\n", err)
	}

	// Check the payload.
	var _p packet.Packet
	copy(_p[:], p)
	payload, err := packet.Payload(&_p)
	if err != nil {
		t.Fatalf("unexpected error getting packet payload: %v\n", err)
	}
	got := payload[0]
	if got != targetPacketNum {
		t.Errorf("payload of found packet is not correct.\nGot: %v, Want: %v\n", got, targetPacketNum)
	}

	// Check the index.
	_got := i / PacketSize
	if _got != targetPacketNum {
		t.Errorf("index of found packet is not correct.\nGot: %v, want: %v\n", _got, targetPacketNum)
	}
}