From 407426419a5d2af2ea921f1c0ae73556b3cb0320 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 20 Mar 2022 18:07:51 +0100
Subject: [PATCH] fix tests

---
 internal/hls/muxer_ts_generator.go | 23 +++++++++++++++--------
 internal/hls/muxer_ts_segment.go   | 11 +++++------
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/internal/hls/muxer_ts_generator.go b/internal/hls/muxer_ts_generator.go
index 1f8e2a13..1ab9a22a 100644
--- a/internal/hls/muxer_ts_generator.go
+++ b/internal/hls/muxer_ts_generator.go
@@ -95,6 +95,7 @@ func newMuxerTSGenerator(
 }
 
 func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
+	now := time.Now()
 	idrPresent := idrPresent(nalus)
 
 	if m.currentSegment == nil {
@@ -104,8 +105,9 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 		}
 
 		// create first segment
-		m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData)
-		m.startPCR = time.Now()
+		m.startPCR = now
+		m.currentSegment = newMuxerTSSegment(now, m.hlsSegmentMaxSize,
+			m.videoTrack, m.writer.WriteData)
 		m.startPTS = pts
 		m.videoDTSEst = h264.NewDTSEstimator()
 		pts = pcrOffset
@@ -118,7 +120,8 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 			(pts-*m.currentSegment.startPTS) >= m.hlsSegmentDuration {
 			m.currentSegment.endPTS = pts
 			m.streamPlaylist.pushSegment(m.currentSegment)
-			m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData)
+			m.currentSegment = newMuxerTSSegment(now, m.hlsSegmentMaxSize,
+				m.videoTrack, m.writer.WriteData)
 		}
 	}
 
@@ -153,7 +156,7 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 		return err
 	}
 
-	err = m.currentSegment.writeH264(m.startPCR, dts, pts, idrPresent, enc)
+	err = m.currentSegment.writeH264(now.Sub(m.startPCR), dts, pts, idrPresent, enc)
 	if err != nil {
 		if m.currentSegment.buf.Len() > 0 {
 			m.streamPlaylist.pushSegment(m.currentSegment)
@@ -166,11 +169,14 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 }
 
 func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error {
+	now := time.Now()
+
 	if m.videoTrack == nil {
 		if m.currentSegment == nil {
 			// create first segment
-			m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData)
-			m.startPCR = time.Now()
+			m.startPCR = now
+			m.currentSegment = newMuxerTSSegment(now, m.hlsSegmentMaxSize,
+				m.videoTrack, m.writer.WriteData)
 			m.startPTS = pts
 			pts = pcrOffset
 		} else {
@@ -182,7 +188,8 @@ func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error {
 				(pts-*m.currentSegment.startPTS) >= m.hlsSegmentDuration {
 				m.currentSegment.endPTS = pts
 				m.streamPlaylist.pushSegment(m.currentSegment)
-				m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData)
+				m.currentSegment = newMuxerTSSegment(now, m.hlsSegmentMaxSize,
+					m.videoTrack, m.writer.WriteData)
 			}
 		}
 	} else {
@@ -210,7 +217,7 @@ func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error {
 		return err
 	}
 
-	err = m.currentSegment.writeAAC(m.startPCR, pts, enc, len(aus))
+	err = m.currentSegment.writeAAC(now.Sub(m.startPCR), pts, enc, len(aus))
 	if err != nil {
 		if m.currentSegment.buf.Len() > 0 {
 			m.streamPlaylist.pushSegment(m.currentSegment)
diff --git a/internal/hls/muxer_ts_segment.go b/internal/hls/muxer_ts_segment.go
index a3e1e41c..56671872 100644
--- a/internal/hls/muxer_ts_segment.go
+++ b/internal/hls/muxer_ts_segment.go
@@ -26,12 +26,11 @@ type muxerTSSegment struct {
 }
 
 func newMuxerTSSegment(
+	now time.Time,
 	hlsSegmentMaxSize uint64,
 	videoTrack *gortsplib.TrackH264,
 	writeData func(*astits.MuxerData) (int, error),
 ) *muxerTSSegment {
-	now := time.Now()
-
 	t := &muxerTSSegment{
 		hlsSegmentMaxSize: hlsSegmentMaxSize,
 		videoTrack:        videoTrack,
@@ -65,7 +64,7 @@ func (t *muxerTSSegment) reader() io.Reader {
 }
 
 func (t *muxerTSSegment) writeH264(
-	startPCR time.Time,
+	pcr time.Duration,
 	dts time.Duration,
 	pts time.Duration,
 	idrPresent bool,
@@ -83,7 +82,7 @@ func (t *muxerTSSegment) writeH264(
 			af = &astits.PacketAdaptationField{}
 		}
 		af.HasPCR = true
-		af.PCR = &astits.ClockReference{Base: int64(time.Since(startPCR).Seconds() * 90000)}
+		af.PCR = &astits.ClockReference{Base: int64(pcr.Seconds() * 90000)}
 		t.pcrSendCounter = 3
 	}
 	t.pcrSendCounter--
@@ -128,7 +127,7 @@ func (t *muxerTSSegment) writeH264(
 }
 
 func (t *muxerTSSegment) writeAAC(
-	startPCR time.Time,
+	pcr time.Duration,
 	pts time.Duration,
 	enc []byte,
 	ausLen int) error {
@@ -140,7 +139,7 @@ func (t *muxerTSSegment) writeAAC(
 		// send PCR once in a while
 		if t.pcrSendCounter == 0 {
 			af.HasPCR = true
-			af.PCR = &astits.ClockReference{Base: int64(time.Since(startPCR).Seconds() * 90000)}
+			af.PCR = &astits.ClockReference{Base: int64(pcr.Seconds() * 90000)}
 			t.pcrSendCounter = 3
 		}
 	}