av/container/mts/payload_test.go

138 lines
3.4 KiB
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 = append(want, 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) != len(got) {
t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got), len(want))
}
// Check frames individually.
for i, frame := range want {
if !reflect.DeepEqual(frame, got[i]) {
t.Fatalf("did not get expected result.\nGot: %v\n, Want: %v\n", got[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
}