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