From 5179cdaceb53831136d676754e9d63788ecca380 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Fri, 10 May 2019 17:38:22 +0930
Subject: [PATCH] container/mts/mpegts_test.go: wrote some helper functions for
 writing psi and frames and cleaned up commenting

---
 container/mts/mpegts.go      |   3 +-
 container/mts/mpegts_test.go | 142 +++++++++++++++++++++--------------
 2 files changed, 86 insertions(+), 59 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index 4d8ec4a8..ef2e6306 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -185,7 +185,7 @@ func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 	for i = 0; i < len(d); i += PacketSize {
 		p := (uint16(d[i+1]&0x1f) << 8) | uint16(d[i+2])
 		if p == pid {
-			pkt = d[i+4 : i+PacketSize]
+			pkt = d[i : i+PacketSize]
 			return
 		}
 	}
@@ -330,6 +330,7 @@ func GetPTSRange(clip []byte) (pts [2]uint64, err error) {
 	copy(_pkt[:], pkt)
 	payload, err := packet.Payload(&_pkt)
 	if err != nil {
+		fmt.Printf("_pkt: %v\n", _pkt)
 		return [2]uint64{}, err
 	}
 
diff --git a/container/mts/mpegts_test.go b/container/mts/mpegts_test.go
index 7c096444..3bd7bf59 100644
--- a/container/mts/mpegts_test.go
+++ b/container/mts/mpegts_test.go
@@ -10,85 +10,44 @@ import (
 	"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
-		maxFrame    = 1000
-		minFrame    = 100
-		rate        = 25 // fps
-		interval    = float64(1) / rate
-		ptsFreq     = 90000 // Hz
+		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(maxFrame-minFrame) + minFrame
+		size := rand.Intn(maxFrameSize-minFrameSize) + minFrameSize
 		frames[i] = make([]byte, size)
 	}
 
 	var clip bytes.Buffer
 
-	// Write PAT and PMT.
-	pat := Packet{
-		PUSI:    true,
-		PID:     PatPid,
-		CC:      0,
-		AFC:     HasPayload,
-		Payload: psi.AddPadding(patTable),
-	}
-	_, err := clip.Write(pat.Bytes(nil))
+	// Write the PSI first.
+	err := writePSI(clip)
 	if err != nil {
-		t.Fatalf("did not expect error writing PAT: %v", err)
-	}
-
-	// Create mts packet from pmt table.
-	pmt := Packet{
-		PUSI:    true,
-		PID:     PmtPid,
-		CC:      0,
-		AFC:     HasPayload,
-		Payload: psi.AddPadding(pmtTable),
-	}
-	_, err = clip.Write(pmt.Bytes(nil))
-	if err != nil {
-		t.Fatalf("did not expected error writing PMT: %v", err)
+		t.Fatalf("did not expect error writing psi: %v", err)
 	}
 
+	// Now write frames.
 	var curTime float64
 	for _, frame := range frames {
 		nextPTS := curTime * ptsFreq
-		// Prepare PES data.
-		pesPkt := pes.Packet{
-			StreamID:     videoStreamID,
-			PDI:          hasPTS,
-			PTS:          uint64(nextPTS),
-			Data:         frame,
-			HeaderLength: 5,
+
+		err = writeFrame(clip, frame, uint64(nextPTS))
+		if err != nil {
+			t.Fatalf("did not expect error writing frame: %v", err)
 		}
 
-		buf := pesPkt.Bytes(nil)
-
-		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 := clip.Write(pkt.Bytes(nil))
-			if err != nil {
-				t.Fatalf("did not expect error writing video packet: %v", err)
-			}
-		}
 		curTime += interval
 	}
 
@@ -102,3 +61,70 @@ func TestGetPTSRange(t *testing.T) {
 		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
+}