container/mts/mpegts.go: FindPSI now returns a map[uint16]uint8 containing program PIDs and their types instead of []psi.PmtElementaryStream

container/mts/mpegts.go

@@ -225,9 +225,8 @@ var (
 )
 
 // FindPSI finds the index of a PAT in an a slice of MPEG-TS and returns, along
-// with a map of meta from the PMT and the PmtElementaryStreams, which contain
+// with a map of meta from the PMT and the stream PIDs and their types.
-// the media PIDs and their types.
+func FindPSI(d []byte) (int, map[uint16]uint8, map[string]string, error) {
-func FindPSI(d []byte) (int, []gotspsi.PmtElementaryStream, map[string]string, error) {
 	if len(d) < PacketSize {
 		return -1, nil, nil, ErrInvalidLen
 	}
@@ -278,7 +277,12 @@ func FindPSI(d []byte) (int, []gotspsi.PmtElementaryStream, map[string]string, e
 		return i, nil, meta, errors.Wrap(err, "could not get streams from PMT")
 	}
 
-	return i, streams, meta, nil
+	var streamMap map[uint16]uint8
+	for _, s := range streams {
+		streamMap[s.ElementaryPid()] = s.StreamType()
+	}
+
+	return i, streamMap, meta, nil
 }
 
 // FillPayload takes a channel and fills the packets Payload field until the

container/mts/mpegts_test.go

@@ -734,9 +734,15 @@ func TestFindPSI(t *testing.T) {
 				t.Fatalf("gotStreams should not be 0 length")
 			}
 
-			s := gotStreams[0]
+			var (
-			gotStreamType := s.StreamType()
+				gotStreamPID  uint16
-			gotStreamPID := s.ElementaryPid()
+				gotStreamType uint8
+			)
+
+			for k, v := range gotStreams {
+				gotStreamPID = k
+				gotStreamType = v
+			}
 
 			if gotStreamType != test.want.streamType {
 				t.Errorf("did not get expected stream type for test %d\nGot: %v\nWant: %v\n", i, gotStreamType, test.want.streamType)