av/container/mts/payload_test.go

package mts

import (
	"bytes"
	"math/rand"
	"reflect"
	"strconv"
	"testing"
	"time"

	"bitbucket.org/ausocean/av/container/mts/meta"
	"bitbucket.org/ausocean/av/container/mts/psi"
)

// TestExtract checks that we can coorectly extract media, pts, id and meta from
// an mpegts stream using Extract.
func TestExtract(t *testing.T) {
	Meta = meta.New()

	const (
		psiInterval  = 5                 // Write PSI at start and after every 5 frames.
		numOfFrames  = 30                // Total number of frames to write.
		maxFrameSize = 1000              // Max frame size to randomly generate.
		minFrameSize = 100               // Min frame size to randomly generate.
		rate         = 25                // Framerate (fps)
		interval     = float64(1) / rate // Time interval between frames.
		ptsFreq      = 90000             // Standard PTS frequency base.
	)

	// Generate randomly sized data for each frame and fill.
	rand.Seed(time.Now().UnixNano())
	frames := make([][]byte, numOfFrames)
	for i := range frames {
		size := rand.Intn(maxFrameSize-minFrameSize) + minFrameSize
		frames[i] = make([]byte, size)
		for j := 0; j < len(frames[i]); j++ {
			frames[i][j] = byte(j)
		}
	}

	var (
		clip bytes.Buffer // This will hold the MPEG-TS data.
		want Clip         // This is the Clip that we should get.
		err  error
	)

	// Now write frames.
	var curTime float64
	for i, frame := range frames {
		// Check to see if it's time to write another lot of PSI.
		if i%psiInterval == 0 && i != len(frames)-1 {
			// We'll add the frame number as meta.
			Meta.Add("frameNum", strconv.Itoa(i))

			err = writePSIWithMeta(&clip)
			if err != nil {
				t.Fatalf("did not expect error writing psi: %v", err)
			}
		}
		nextPTS := uint64(curTime * ptsFreq)

		err = writeFrame(&clip, frame, uint64(nextPTS))
		if err != nil {
			t.Fatalf("did not expect error writing frame: %v", err)
		}

		curTime += interval

		// Need the meta map for the new expected Frame.
		metaMap, err := meta.GetAllAsMap(Meta.Encode())
		if err != nil {
			t.Fatalf("did not expect error getting meta map: %v", err)
		}

		// Create an equivalent Frame and append to our Clip want.
		want.frames = append(want.frames, Frame{
			Media: frame,
			PTS:   nextPTS,
			ID:    H264ID,
			Meta:  metaMap,
		})
	}

	// Now use Extract to get frames from clip.
	got, err := Extract(clip.Bytes())
	if err != nil {
		t.Fatalf("did not expect error using Extract. Err: %v", err)
	}

	// Check length of got and want.
	if len(want.frames) != len(got.frames) {
		t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got.frames), len(want.frames))
	}

	// Check frames individually.
	for i, frame := range want.frames {
		if !reflect.DeepEqual(frame, got.frames[i]) {
			t.Fatalf("did not get expected result.\nGot: %v\n, Want: %v\n", got.frames[i], frame)
		}
	}
}

// writePSIWithMeta writes PSI to b with updated metadata.
func writePSIWithMeta(b *bytes.Buffer) error {
	// Write PAT.
	pat := Packet{
		PUSI:    true,
		PID:     PatPid,
		CC:      0,
		AFC:     HasPayload,
		Payload: psi.AddPadding(patTable),
	}
	_, err := b.Write(pat.Bytes(nil))
	if err != nil {
		return err
	}

	// Update the meta in the pmt table.
	pmtTable, err = updateMeta(pmtTable)
	if err != nil {
		return err
	}

	// Write PMT.
	pmt := Packet{
		PUSI:    true,
		PID:     PmtPid,
		CC:      0,
		AFC:     HasPayload,
		Payload: psi.AddPadding(pmtTable),
	}
	_, err = b.Write(pmt.Bytes(nil))
	if err != nil {
		return err
	}
	return nil
}