container/mts: wrote test for Clip.Bytes and generalised logic for generating frames into genFrames func

container/mts/payload_test.go

@@ -27,16 +27,7 @@ func TestExtract(t *testing.T) {
 		ptsFreq      = 90000             // Standard PTS frequency base.
 	)
 
-	// Generate randomly sized data for each frame and fill.
-	rand.Seed(time.Now().UnixNano())
-	frames := make([][]byte, numOfFrames)
-	for i := range frames {
-		size := rand.Intn(maxFrameSize-minFrameSize) + minFrameSize
-		frames[i] = make([]byte, size)
-		for j := 0; j < len(frames[i]); j++ {
-			frames[i][j] = byte(j)
-		}
-	}
+	frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)
 
 	var (
 		clip bytes.Buffer // This will hold the MPEG-TS data.
@@ -135,3 +126,84 @@ func writePSIWithMeta(b *bytes.Buffer) error {
 	}
 	return nil
 }
+
+// TestClipBytes checks that Clip.Bytes correctly returns the concatendated media
+// data from the Clip's frames slice.
+func TestClipBytes(t *testing.T) {
+	Meta = meta.New()
+
+	const (
+		psiInterval  = 5                 // Write PSI at start and after every 5 frames.
+		numOfFrames  = 30                // Total number of frames to write.
+		maxFrameSize = 1000              // Max frame size to randomly generate.
+		minFrameSize = 100               // Min frame size to randomly generate.
+		rate         = 25                // Framerate (fps)
+		interval     = float64(1) / rate // Time interval between frames.
+		ptsFreq      = 90000             // Standard PTS frequency base.
+	)
+
+	frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)
+
+	var (
+		clip bytes.Buffer // This will hold the MPEG-TS data.
+		want []byte       // This is the Clip that we should get.
+		err  error
+	)
+
+	// Now write frames.
+	var curTime float64
+	for i, frame := range frames {
+		// Check to see if it's time to write another lot of PSI.
+		if i%psiInterval == 0 && i != len(frames)-1 {
+			// We'll add the frame number as meta.
+			Meta.Add("frameNum", strconv.Itoa(i))
+
+			err = writePSIWithMeta(&clip)
+			if err != nil {
+				t.Fatalf("did not expect error writing psi: %v", err)
+			}
+		}
+		nextPTS := uint64(curTime * ptsFreq)
+
+		err = writeFrame(&clip, frame, uint64(nextPTS))
+		if err != nil {
+			t.Fatalf("did not expect error writing frame: %v", err)
+		}
+
+		curTime += interval
+
+		// Append the frame straight to the expected pure media slice.
+		want = append(want, frame...)
+	}
+
+	// Now use Extract to get Clip and then use Bytes to get the slice of straight media.
+	gotClip, err := Extract(clip.Bytes())
+	if err != nil {
+		t.Fatalf("did not expect error using Extract. Err: %v", err)
+	}
+	got := gotClip.Bytes()
+
+	// Check length and equality of got and want.
+	if len(want) != len(got) {
+		t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got), len(want))
+	}
+	if !bytes.Equal(want, got) {
+		t.Error("did not get expected result")
+	}
+}
+
+// genFrames is a helper function to generate a series of dummy media frames
+// with randomized size. n is the number of frames to generate, min is the min
+// size is min size of random frame and max is max size of random frames.
+func genFrames(n, min, max int) [][]byte {
+	// Generate randomly sized data for each frame and fill.
+	rand.Seed(time.Now().UnixNano())
+	frames := make([][]byte, n)
+	for i := range frames {
+		frames[i] = make([]byte, rand.Intn(max-min)+min)
+		for j := 0; j < len(frames[i]); j++ {
+			frames[i][j] = byte(j)
+		}
+	}
+	return frames
+}