av/container/mts/audio_test.go

/*
NAME
  audio_test.go

DESCRIPTION
  See Readme.md

AUTHOR
  Trek Hopton <trek@ausocean.org>

LICENSE
  audio_test.go is Copyright (C) 2017-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
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package mts

import (
	"bytes"
	"io/ioutil"
	"log"
	"testing"

	"bitbucket.org/ausocean/av/container/mts/meta"
	"github.com/Comcast/gots/packet"
	"github.com/Comcast/gots/pes"
)

// TestEncodePcm tests the mpegts encoder's ability to encode pcm audio data.
// It reads and encodes input pcm data into mpegts, then decodes the mpegts compares the result to the input pcm.
func TestEncodePcm(t *testing.T) {
	Meta = meta.New()

	var b []byte
	buf := bytes.NewBuffer(b)
	sampleRate := 48000
	sampleSize := 2
	writeSize := 16000
	writeFreq := float64(sampleRate*sampleSize) / float64(writeSize)
	e := NewEncoder(buf, writeFreq, Audio)

	inPath := "../../../test/test-data/av/input/sweep_400Hz_20000Hz_-3dBFS_5s_48khz.pcm"
	inPcm, err := ioutil.ReadFile(inPath)
	if err != nil {
		log.Fatal(err)
	}

	// Encode pcm to mts and get the resulting bytes.
	_, err = e.Write(inPcm)
	if err != nil {
		log.Fatal(err)
	}
	clip := buf.Bytes()

	// Decode the mts packets to extract the original data
	var pkt packet.Packet
	pesPacket := make([]byte, 0, writeSize)
	got := make([]byte, 0, len(inPcm))
	i := 0
	for i+PacketSize <= len(clip) {
		copy(pkt[:], clip[i:i+PacketSize])
		if pkt.PID() == audioPid {
			if pkt.PayloadUnitStartIndicator() {
				payload, err := pkt.Payload()
				if err != nil {
					t.Fatalf("Unexpected err: %v\n", err)
				}
				pesPacket = append(pesPacket, payload...)
				i += PacketSize
				first := true
				for (first || !pkt.PayloadUnitStartIndicator()) && i+PacketSize <= len(clip) {
					first = false
					copy(pkt[:], clip[i:i+PacketSize])
					payload, err := pkt.Payload()
					if err != nil {
						t.Fatalf("Unexpected err: %v\n", err)
					}

					pesPacket = append(pesPacket, payload...)
					i += PacketSize
				}
			}
			pesHeader, err := pes.NewPESHeader(pesPacket)
			if err != nil {
				t.Fatalf("Unexpected err: %v\n", err)
			}
			got = append(got, pesHeader.Data()...)
		} else {
			i += PacketSize
		}
	}

	// Compare encoded data with original data.
	if !bytes.Equal(got, inPcm) {
		t.Error("Error, unexpected output")
	}
}