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 {
		// Check media data.
		wantMedia := frame.Media
		gotMedia := got.frames[i].Media
		if !bytes.Equal(wantMedia, gotMedia) {
			t.Fatalf("did not get expected data for frame: %v\nGot: %v\nWant: %v\n", i, gotMedia, wantMedia)
		}

		// Check stream ID.
		wantID := frame.ID
		gotID := got.frames[i].ID
		if wantID != gotID {
			t.Fatalf("did not get expected ID for frame: %v\nGot: %v\nWant: %v\n", i, gotID, wantID)
		}

		// Check meta.
		wantMeta := frame.Meta
		gotMeta := got.frames[i].Meta
		if !reflect.DeepEqual(wantMeta, gotMeta) {
			t.Fatalf("did not get expected meta for frame: %v\nGot: %v\nwant: %v\n", i, gotMeta, wantMeta)
		}
	}
}

// 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
}

// TestBytesForPTSInterval checks that BytesForPTSInterval can correctly return
// a slice of media data corresponding to a given PTS interval.
func TestBytesForPTSInterval(t *testing.T) {
	const (
		numOfTestFrames = 10
		ptsInterval     = 4
		frameSize       = 3
	)

	clip := &Clip{}

	// Generate test frames.
	for i := 0; i < numOfTestFrames; i++ {
		clip.backing = append(clip.backing, []byte{byte(i), byte(i), byte(i)}...)
		clip.frames = append(
			clip.frames,
			Frame{
				Media: clip.backing[i*frameSize : (i+1)*frameSize],
				PTS:   uint64(i * ptsInterval),
				idx:   i * frameSize,
			},
		)
	}

	// We test each of these scenarios.
	tests := []struct {
		from   uint64
		to     uint64
		expect []byte
		err    error
	}{
		{
			from: 6,
			to:   15,
			expect: []byte{
				0x01, 0x01, 0x01,
				0x02, 0x02, 0x02,
				0x03, 0x03, 0x03,
			},
			err: nil,
		},
		{
			from: 4,
			to:   16,
			expect: []byte{
				0x01, 0x01, 0x01,
				0x02, 0x02, 0x02,
				0x03, 0x03, 0x03,
			},
			err: nil,
		},
		{
			from:   10,
			to:     5,
			expect: nil,
			err:    errPTSRange,
		},
		{
			from:   50,
			to:     70,
			expect: nil,
			err:    errPTSLowerBound,
		},
		{
			from:   5,
			to:     70,
			expect: nil,
			err:    errPTSUpperBound,
		},
	}

	// Run tests.
	for i, test := range tests {
		got, err := clip.BytesForPTSInterval(test.from, test.to)

		// First check the error.
		if err != nil && err != test.err {
			t.Errorf("unexpected error: %v for test: %v from BytesForPTSInterval", err, i)
			continue
		} else if err != test.err {
			t.Errorf("expected to get error: %v for test: %v from BytesForPTSInterval", test.err, i)
			continue
		}

		// Now check data.
		if test.err == nil && !bytes.Equal(test.expect, got) {
			t.Errorf("did not get expected data for test: %v from BytesForPTSInterval.\n Got: %v\n, Want: %v\n", i, got, test.expect)
		}
	}
}