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.
	)

	frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)

	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
}

// TestClipBytes checks that Clip.Bytes correctly returns the concatendated media
// data from the Clip's frames slice.
func TestClipBytes(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.
	)

	frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)

	var (
		clip bytes.Buffer // This will hold the MPEG-TS data.
		want []byte       // 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

		// Append the frame straight to the expected pure media slice.
		want = append(want, frame...)
	}

	// Now use Extract to get Clip and then use Bytes to get the slice of straight media.
	gotClip, err := Extract(clip.Bytes())
	if err != nil {
		t.Fatalf("did not expect error using Extract. Err: %v", err)
	}
	got := gotClip.Bytes()

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

// genFrames is a helper function to generate a series of dummy media frames
// with randomized size. n is the number of frames to generate, min is the min
// size is min size of random frame and max is max size of random frames.
func genFrames(n, min, max int) [][]byte {
	// Generate randomly sized data for each frame and fill.
	rand.Seed(time.Now().UnixNano())
	frames := make([][]byte, n)
	for i := range frames {
		frames[i] = make([]byte, rand.Intn(max-min)+min)
		for j := 0; j < len(frames[i]); j++ {
			frames[i][j] = byte(j)
		}
	}
	return frames
}