Added test for MPEG-TS passthrough writing to recorder.test.go

Author: zach-grant-skydio

internal/recorder/recorder_test.go

@@ -1,12 +1,15 @@
 package recorder
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
 	"time"
 
+	"github.com/pion/rtp"
+
 	"github.com/bluenviron/gortsplib/v4/pkg/description"
 	rtspformat "github.com/bluenviron/gortsplib/v4/pkg/format"
 	"github.com/bluenviron/mediacommon/v2/pkg/codecs/h265"
@@ -539,6 +542,159 @@ func TestRecorderSkipTracksFull(t *testing.T) {
 	}
 }
 
+func TestRecorderMPEGTSPassthrough(t *testing.T) {
+	// Create a test MPEG-TS packet
+	testData := []byte{
+		0x47, 0x40, 0x00, 0x10, 0x00, // TS header with PID 0x1000 (video)
+		// Add some dummy payload - this is a PAT (Program Association Table)
+		0x00, 0x00, 0xB0, 0x0D, 0x00, 0x00, 0xC1, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+	}
+
+	t.Run("basic", func(t *testing.T) {
+		// Create a temporary directory for the test
+		dir, err := os.MkdirTemp("", "mediamtx-agent")
+		require.NoError(t, err)
+		defer os.RemoveAll(dir)
+
+		// Create the output directory
+		recordPath := filepath.Join(dir, "%path/%Y-%m-%d_%H-%M-%S-%f")
+
+		// Create a stream with MPEG-TS format
+		mpegtsFormat := &rtspformat.MPEGTS{}
+
+		desc := &description.Session{Medias: []*description.Media{
+			{
+				Type:    description.MediaTypeVideo,
+				Formats: []rtspformat.Format{mpegtsFormat},
+			},
+		}}
+
+		strm := &stream.Stream{
+			WriteQueueSize:     512,
+			UDPMaxPayloadSize:  1472,
+			Desc:               desc,
+			GenerateRTPPackets: false, // Disable RTP packet generation
+			Parent:             test.NilLogger,
+		}
+		err = strm.Initialize()
+		require.NoError(t, err)
+		defer strm.Close()
+
+		// Create a recorder with MPEG-TS format
+		r := &Recorder{
+			PathFormat:      recordPath,
+			Format:          conf.RecordFormatMPEGTS,
+			PartDuration:    1 * time.Second,
+			SegmentDuration: 2 * time.Second,
+			PathName:        "mypath",
+			Stream:          strm,
+			Parent:          test.NilLogger,
+		}
+
+		r.Initialize()
+
+		// Verify the recorder instance was created
+		require.NotNil(t, r.currentInstance, "Recorder instance is nil")
+
+		// Create a channel to receive data from the stream
+		dataChan := make(chan []byte, 10)
+
+		strm.AddReader(r, desc.Medias[0], mpegtsFormat, func(u unit.Unit) error {
+			rtpPackets := u.GetRTPPackets()
+
+			for _, pkt := range rtpPackets {
+				if pkt == nil {
+					continue
+				}
+
+				// Send the payload to the channel
+				if len(pkt.Payload) > 0 {
+					dataChan <- pkt.Payload
+				}
+			}
+
+			return nil
+		})
+
+		strm.StartReader(r)
+		strm.WaitRunningReader()
+
+		// Write multiple RTP packets to ensure we have enough data
+		for i := 0; i < 10; i++ {
+			// Create a new RTP packet for each iteration with updated timestamp and sequence number
+			pkt := &rtp.Packet{
+				Header: rtp.Header{
+					Version:        2,
+					PayloadType:    33, // Standard MPEG-TS payload type (RFC 2250)
+					SequenceNumber: 100 + uint16(i),
+					Timestamp:      123456 + uint32(i*90000), // 1 second apart
+					SSRC:           0x9D8F,
+				},
+				Payload: testData,
+			}
+
+			// Use the MPEG-TS format directly
+			internalFormat := mpegtsFormat
+
+			// Write the RTP packet to the stream
+			strm.WriteRTPPacket(desc.Medias[0], internalFormat, pkt, time.Now(), 0)
+		}
+
+		// Check if the recorder instance was properly initialized
+		if r.currentInstance == nil {
+			t.Fatal("Recorder currentInstance is nil")
+		}
+
+		// Manually flush the buffer writer to ensure all data is written
+		if r.currentInstance != nil && r.currentInstance.format2 != nil {
+			if mpegtsFormat, ok := r.currentInstance.format2.(*formatMPEGTS); ok && mpegtsFormat.bw != nil {
+				err = mpegtsFormat.bw.Flush()
+				require.NoError(t, err, "Failed to flush buffer writer")
+
+				// Close the current segment if it exists
+				if mpegtsFormat.currentSegment != nil {
+					err = mpegtsFormat.currentSegment.close()
+					require.NoError(t, err, "Failed to close segment")
+					mpegtsFormat.currentSegment = nil
+				}
+			}
+		}
+
+		// Close the recorder to flush any remaining data
+		r.Close()
+
+		// Verify the recording was created
+		entries, err := os.ReadDir(filepath.Join(dir, "mypath"))
+		require.NoError(t, err)
+
+		// Find the .ts file
+		var tsFile string
+		for _, entry := range entries {
+			if filepath.Ext(entry.Name()) == ".ts" {
+				tsFile = filepath.Join(dir, "mypath", entry.Name())
+				break
+			}
+		}
+
+		if tsFile == "" {
+			t.Fatalf("No .ts file found in %s. Directory contents: %v", recordPath, entries)
+		}
+
+		// Verify the file contains our test data
+		data, err := os.ReadFile(tsFile)
+		require.NoError(t, err)
+		require.Greater(t, len(data), 0, "Recorded file is empty")
+		require.True(t, bytes.Contains(data, testData), "Test data not found in output")
+	})
+}
+
 func TestRecorderFMP4SegmentSwitch(t *testing.T) {
 	desc := &description.Session{Medias: []*description.Media{
 		{