av/container/mts/mpegts_test.go

package mts

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

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

// TestGetPTSRange checks that GetPTSRange can correctly get the first and last
// PTS in an MPEGTS clip.
func TestGetPTSRange(t *testing.T) {
	const (
		numOfFrames  = 20
		maxFrameSize = 1000
		minFrameSize = 100
		rate         = 25                // fps
		interval     = float64(1) / rate // s
		ptsFreq      = 90000             // Hz
	)

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

	var clip bytes.Buffer

	// Write the PSI first.
	err := writePSI(clip)
	if err != nil {
		t.Fatalf("did not expect error writing psi: %v", err)
	}

	// Now write frames.
	var curTime float64
	for _, frame := range frames {
		nextPTS := curTime * ptsFreq

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

		curTime += interval
	}

	got, err := GetPTSRange(clip.Bytes())
	if err != nil {
		t.Fatalf("did not expect error getting PTS range: %v", err)
	}

	want := [2]uint64{0, uint64((numOfFrames - 1) * interval * ptsFreq)}
	if got != want {
		t.Errorf("did not get expected result.\n Got: %v\n Want: %v\n", got, want)
	}
}

// writePSI is a helper function write the PSI found at the start of a clip.
func writePSI(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
	}

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

// writeFrame is a helper function used to form a PES packet from a frame, and
// then fragment this across MPEGTS packets where they are then written to the
// given buffer.
func writeFrame(b bytes.Buffer, frame []byte, pts uint64) error {
	// Prepare PES data.
	pesPkt := pes.Packet{
		StreamID:     videoStreamID,
		PDI:          hasPTS,
		PTS:          pts,
		Data:         frame,
		HeaderLength: 5,
	}
	buf := pesPkt.Bytes(nil)

	// Write PES data acroos MPEGTS packets.
	pusi := true
	for len(buf) != 0 {
		pkt := Packet{
			PUSI: pusi,
			PID:  videoPid,
			RAI:  pusi,
			CC:   0,
			AFC:  hasAdaptationField | hasPayload,
			PCRF: pusi,
		}
		n := pkt.FillPayload(buf)
		buf = buf[n:]

		pusi = false
		_, err := b.Write(pkt.Bytes(nil))
		if err != nil {
			return err
		}
	}
	return nil
}