mirror of https://bitbucket.org/ausocean/av.git
138 lines
3.4 KiB
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
|
|
}
|