From a8aec484dfc22b1ada94ac779d24df279d490ae4 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Tue, 7 May 2019 02:20:07 +0930
Subject: [PATCH] container/mts: wrote test TestFindPid and corrected bug

Completed test TestFindPid to validate function of FindPid func in mpegts.go. Through this
process, it was found that there was a bug in this func, whereby the returned packet was
not complete due to indexing issues. The bug is fixed and the test passes.
---
 container/mts/mpegts.go      |  4 ++--
 container/mts/mpegts_test.go | 41 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index ed8cbe02..af8cab8d 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -179,12 +179,12 @@ func FindPat(d []byte) ([]byte, int, error) {
 // is found, then it is returned along with its index, otherwise nil, -1 and an error is returned.
 func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 	if len(d) < PacketSize {
-		return nil, -1, errors.New("Mmpegts data not of valid length")
+		return nil, -1, errors.New("Mpegts data not of valid length")
 	}
 	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
 		}
 	}
diff --git a/container/mts/mpegts_test.go b/container/mts/mpegts_test.go
index 92d165ca..8322df72 100644
--- a/container/mts/mpegts_test.go
+++ b/container/mts/mpegts_test.go
@@ -30,6 +30,8 @@ package mts
 import (
 	"bytes"
 	"testing"
+
+	"github.com/Comcast/gots/packet"
 )
 
 func TestBytes(t *testing.T) {
@@ -94,3 +96,42 @@ func TestBytes(t *testing.T) {
 		}
 	}
 }
+
+func TestFindPid(t *testing.T) {
+	const targetPacketNum, numOfPackets, targetPid, stdPid = 6, 15, 1, 0
+
+	var stream []byte
+	for i := 0; i < numOfPackets; i++ {
+		pid := uint16(0)
+		if i == targetPacketNum {
+			pid = 1
+		}
+		p := Packet{
+			PID: pid,
+			AFC: hasPayload | hasAdaptationField,
+		}
+		p.FillPayload([]byte{byte(i)})
+		stream = append(stream, p.Bytes(nil)...)
+	}
+
+	p, i, err := FindPid(stream, targetPid)
+	if err != nil {
+		t.Fatalf("unexpected error finding PID: %v\n", err)
+	}
+
+	var _p packet.Packet
+	copy(_p[:], p)
+	payload, err := packet.Payload(&_p)
+	if err != nil {
+		t.Fatalf("unexpected error getting packet payload: %v\n", err)
+	}
+	got := payload[0]
+	if got != targetPacketNum {
+		t.Errorf("payload of found packet is not correct.\nGot: %v, Want: %v\n", got, targetPacketNum)
+	}
+
+	_got := i / PacketSize
+	if _got != targetPacketNum {
+		t.Errorf("index of found packet is not correct.\nGot: %v, want: %v\n", _got, targetPacketNum)
+	}
+}