Add TsSegmenter class

- Multi-file segmenter implementation.
- The samples are passed to PesPacketGenerator to generate PesPackets.
  PesPackets are passed to TsWriter to write them to file.

Issue #84

Change-Id: Ia4cae2abe2e7df46016dcdb791baaab2818aea03

packager/media/formats/mp2t/mp2t.gyp

@@ -22,10 +22,10 @@
         'es_parser.h',
         'mp2t_media_parser.cc',
         'mp2t_media_parser.h',
-        'pes_packet_generator.cc',
-        'pes_packet_generator.h',
         'pes_packet.cc',
         'pes_packet.h',
+        'pes_packet_generator.cc',
+        'pes_packet_generator.h',
         'ts_packet.cc',
         'ts_packet.h',
         'ts_section_pat.cc',
@@ -36,6 +36,8 @@
         'ts_section_pmt.h',
         'ts_section_psi.cc',
         'ts_section_psi.h',
+        'ts_segmenter.cc',
+        'ts_segmenter.h',
         'ts_writer.cc',
         'ts_writer.h',
       ],
@@ -54,6 +56,7 @@
         'es_parser_h264_unittest.cc',
         'mp2t_media_parser_unittest.cc',
         'pes_packet_generator_unittest.cc',
+        'ts_segmenter_unittest.cc',
         'ts_writer_unittest.cc',
       ],
       'dependencies': [

packager/media/formats/mp2t/pes_packet_generator.h

@@ -29,38 +29,39 @@ namespace mp2t {
 class PesPacket;
 
 /// Generates PesPackets from MediaSamples.
+/// Methods are virtual for mocking.
 class PesPacketGenerator {
  public:
   PesPacketGenerator();
-  ~PesPacketGenerator();
+  virtual ~PesPacketGenerator();
 
   /// Initialize the object. This clears the internal state first so any
   /// PesPackets that have not been flushed will be lost.
   /// @param stream is the stream info for the elementary stream that will be
   ///        added via PushSample().
   /// @return true on success, false otherwise.
-  bool Initialize(const StreamInfo& stream);
+  virtual bool Initialize(const StreamInfo& stream);
 
   /// Add a sample to the generator. This does not necessarily increase
   /// NumberOfReadyPesPackets().
   /// If this returns false, the object may end up in an undefined state.
   /// @return true on success, false otherwise.
-  bool PushSample(scoped_refptr<MediaSample> sample);
+  virtual bool PushSample(scoped_refptr<MediaSample> sample);
 
   /// @return The number of PES packets that are ready to be consumed.
-  size_t NumberOfReadyPesPackets();
+  virtual size_t NumberOfReadyPesPackets();
 
   /// Removes the next PES packet from the stream and returns it. Must have at
   /// least one packet ready.
   /// @return Next PES packet that is ready.
-  scoped_ptr<PesPacket> GetNextPesPacket();
+  virtual scoped_ptr<PesPacket> GetNextPesPacket();
 
   /// Flush the object. This may create more PesPackets with the stored
   /// samples.
   /// It is safe to call NumberOfReadyPesPackets() and GetNextPesPacket() after
   /// this.
   /// @return true on success, false otherwise.
-  bool Flush();
+  virtual bool Flush();
 
  private:
   friend class PesPacketGeneratorTest;

packager/media/formats/mp2t/ts_segmenter.cc

@@ -0,0 +1,129 @@
+// Copyright 2016 Google Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#include "packager/media/formats/mp2t/ts_segmenter.h"
+
+#include <memory>
+
+#include "packager/media/base/muxer_util.h"
+#include "packager/media/base/status.h"
+
+namespace edash_packager {
+namespace media {
+namespace mp2t {
+
+namespace {
+const int kTsTimescale = 90000;
+}  // namespace
+
+TsSegmenter::TsSegmenter(const MuxerOptions& options)
+    : muxer_options_(options),
+      ts_writer_(new TsWriter()),
+      pes_packet_generator_(new PesPacketGenerator()) {}
+TsSegmenter::~TsSegmenter() {}
+
+Status TsSegmenter::Initialize(const StreamInfo& stream_info) {
+  if (muxer_options_.segment_template.empty())
+    return Status(error::MUXER_FAILURE, "Segment template not specified.");
+  if (!ts_writer_->Initialize(stream_info))
+    return Status(error::MUXER_FAILURE, "Failed to initialize TsWriter.");
+  if (!pes_packet_generator_->Initialize(stream_info)) {
+    return Status(error::MUXER_FAILURE,
+                  "Failed to initialize PesPacketGenerator.");
+  }
+
+  return Status::OK;
+}
+
+Status TsSegmenter::Finalize() {
+  if (!pes_packet_generator_->Flush()) {
+    return Status(error::MUXER_FAILURE,
+                  "Failed to finalize PesPacketGenerator.");
+  }
+
+  Status status = WritePesPacketsToFiles();
+  if (!status.ok())
+    return status;
+
+  if (!ts_writer_file_opened_)
+    return Status::OK;
+
+  if (!ts_writer_->FinalizeSegment())
+    return Status(error::MUXER_FAILURE, "Failed to finalize TsPacketWriter.");
+  ts_writer_file_opened_ = false;
+  return Status::OK;
+}
+
+Status TsSegmenter::AddSample(scoped_refptr<MediaSample> sample) {
+  if (!pes_packet_generator_->PushSample(sample)) {
+    return Status(error::MUXER_FAILURE,
+                  "Failed to add sample to PesPacketGenerator.");
+  }
+  // TODO(rkuriowa): Only segment files before a key frame.
+  return WritePesPacketsToFiles();
+}
+
+void TsSegmenter::InjectTsWriterForTesting(scoped_ptr<TsWriter> writer) {
+  ts_writer_ = writer.Pass();
+}
+
+void TsSegmenter::InjectPesPacketGeneratorForTesting(
+    scoped_ptr<PesPacketGenerator> generator) {
+  pes_packet_generator_ = generator.Pass();
+}
+
+void TsSegmenter::SetTsWriterFileOpenedForTesting(bool value) {
+  ts_writer_file_opened_ = value;
+}
+
+Status TsSegmenter::FinalizeSegmentIfPastSegmentDuration() {
+  if (current_segment_duration_ > muxer_options_.segment_duration) {
+    if (!ts_writer_->FinalizeSegment())
+      return Status(error::FILE_FAILURE, "Failed to finalize segment.");
+    ts_writer_file_opened_ = false;
+    current_segment_duration_ = 0.0;
+  }
+  return Status::OK;
+}
+
+Status TsSegmenter::OpenNewSegmentIfClosed(uint32_t next_pts) {
+  if (ts_writer_file_opened_)
+    return Status::OK;
+  const std::string segment_name =
+      GetSegmentName(muxer_options_.segment_template, next_pts,
+                     segment_number_++, muxer_options_.bandwidth);
+  if (!ts_writer_->NewSegment(segment_name))
+    return Status(error::MUXER_FAILURE, "Failed to initilize TsPacketWriter.");
+  ts_writer_file_opened_ = true;
+  return Status::OK;
+}
+
+Status TsSegmenter::WritePesPacketsToFiles() {
+  while (pes_packet_generator_->NumberOfReadyPesPackets() > 0u) {
+    scoped_ptr<PesPacket> pes_packet =
+        pes_packet_generator_->GetNextPesPacket();
+
+    Status status = OpenNewSegmentIfClosed(pes_packet->pts());
+    if (!status.ok())
+      return status;
+
+    const double pes_packet_duration = pes_packet->duration();
+
+    if (!ts_writer_->AddPesPacket(pes_packet.Pass()))
+      return Status(error::MUXER_FAILURE, "Failed to add PES packet.");
+
+    current_segment_duration_ += pes_packet_duration / kTsTimescale;
+
+    status = FinalizeSegmentIfPastSegmentDuration();
+    if (!status.ok())
+      return status;
+  }
+  return Status::OK;
+}
+
+}  // namespace mp2t
+}  // namespace media
+}  // namespace edash_packager

packager/media/formats/mp2t/ts_segmenter.h

@@ -0,0 +1,84 @@
+// Copyright 2016 Google Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#ifndef PACKAGER_MEDIA_FORMATS_MP2T_TS_SEGMENTER_H_
+#define PACKAGER_MEDIA_FORMATS_MP2T_TS_SEGMENTER_H_
+
+#include "packager/base/memory/scoped_ptr.h"
+#include "packager/media/base/media_stream.h"
+#include "packager/media/base/muxer_options.h"
+#include "packager/media/base/status.h"
+#include "packager/media/file/file.h"
+#include "packager/media/formats/mp2t/pes_packet_generator.h"
+#include "packager/media/formats/mp2t/ts_writer.h"
+
+namespace edash_packager {
+namespace media {
+namespace mp2t {
+
+// TODO(rkuroiwa): For now, this implements multifile segmenter. Like other
+// make this an abstract super class and implement multifile and single file
+// segmenters.
+class TsSegmenter {
+ public:
+  /// @param options is the options for this muxer. This must stay valid
+  ///        throughout the life time of the instance.
+  explicit TsSegmenter(const MuxerOptions& options);
+  ~TsSegmenter();
+
+  /// Initialize the object.
+  /// @param stream_info is the stream info for the segmenter.
+  /// @return OK on success.
+  Status Initialize(const StreamInfo& stream_info);
+
+  /// Finalize the segmenter.
+  /// @return OK on success.
+  Status Finalize();
+
+  /// @param sample gets added to this object.
+  /// @return OK on success.
+  Status AddSample(scoped_refptr<MediaSample> sample);
+
+  /// Only for testing.
+  void InjectTsWriterForTesting(scoped_ptr<TsWriter> writer);
+
+  /// Only for testing.
+  void InjectPesPacketGeneratorForTesting(
+      scoped_ptr<PesPacketGenerator> generator);
+
+  /// Only for testing.
+  void SetTsWriterFileOpenedForTesting(bool value);
+
+ private:
+  Status FinalizeSegmentIfPastSegmentDuration();
+  Status OpenNewSegmentIfClosed(uint32_t next_pts);
+
+  // Writes PES packets (carried in TsPackets) to file(s). This will
+  // segment appropriately. The state of file may be open or closed after
+  // calling this.
+  Status WritePesPacketsToFiles();
+
+  const MuxerOptions& muxer_options_;
+
+  // in seconds.
+  double current_segment_duration_ = 0.0;
+
+  // Used for segment template.
+  uint64_t segment_number_ = 0;
+
+  scoped_ptr<TsWriter> ts_writer_;
+  // Set to true if TsWriter::NewFile() succeeds, set to false after
+  // TsWriter::FinalizeFile() succeeds.
+  bool ts_writer_file_opened_ = false;
+  scoped_ptr<PesPacketGenerator> pes_packet_generator_;
+
+  DISALLOW_COPY_AND_ASSIGN(TsSegmenter);
+};
+
+}  // namespace mp2t
+}  // namespace media
+}  // namespace edash_packager
+#endif  // PACKAGER_MEDIA_FORMATS_MP2T_TS_SEGMENTER_H_

packager/media/formats/mp2t/ts_segmenter_unittest.cc

@@ -0,0 +1,382 @@
+// Copyright 2016 Google Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "packager/media/base/audio_stream_info.h"
+#include "packager/media/base/test/status_test_util.h"
+#include "packager/media/base/video_stream_info.h"
+#include "packager/media/formats/mp2t/ts_segmenter.h"
+
+namespace edash_packager {
+namespace media {
+namespace mp2t {
+
+using ::testing::InSequence;
+using ::testing::Return;
+using ::testing::Sequence;
+using ::testing::StrEq;
+using ::testing::_;
+
+namespace {
+
+// All data here is bogus. They are used to create VideoStreamInfo but the
+// actual values don't matter at all.
+const bool kIsKeyFrame = true;
+const VideoCodec kH264VideoCodec = VideoCodec::kCodecH264;
+const uint8_t kExtraData[] = {
+    0x00,
+};
+const int kTrackId = 0;
+const uint32_t kTimeScale = 90000;
+const uint64_t kDuration = 180000;
+const char kCodecString[] = "avc1";
+const char kLanguage[] = "eng";
+const uint32_t kWidth = 1280;
+const uint32_t kHeight = 720;
+const uint32_t kPixelWidth = 1;
+const uint32_t kPixelHeight = 1;
+const uint16_t kTrickPlayRate = 1;
+const uint8_t kNaluLengthSize = 1;
+const bool kIsEncrypted = false;
+
+class MockPesPacketGenerator : public PesPacketGenerator {
+ public:
+  MOCK_METHOD1(Initialize, bool(const StreamInfo& info));
+  MOCK_METHOD1(PushSample, bool(scoped_refptr<MediaSample> sample));
+  MOCK_METHOD0(NumberOfReadyPesPackets, size_t());
+
+  // Hack found at the URL below for mocking methods that return scoped_ptr.
+  // https://groups.google.com/a/chromium.org/forum/#!topic/chromium-dev/01sDxsJ1OYw
+  MOCK_METHOD0(GetNextPesPacketMock, PesPacket*());
+  scoped_ptr<PesPacket> GetNextPesPacket() override {
+    return scoped_ptr<PesPacket>(GetNextPesPacketMock());
+  }
+
+  MOCK_METHOD0(Flush, bool());
+};
+
+class MockTsWriter : public TsWriter {
+ public:
+  MOCK_METHOD1(Initialize, bool(const StreamInfo& stream_info));
+  MOCK_METHOD1(NewSegment, bool(const std::string& file_name));
+  MOCK_METHOD0(FinalizeSegment, bool());
+
+  // Similar to the hack above but takes a scoped_ptr.
+  MOCK_METHOD1(AddPesPacketMock, bool(PesPacket* pes_packet));
+  bool AddPesPacket(scoped_ptr<PesPacket> pes_packet) override {
+    // No need to keep the pes packet around for the current tests.
+    return AddPesPacketMock(pes_packet.get());
+  }
+
+  MOCK_METHOD0(NotifyReinjectPsi, bool());
+  MOCK_CONST_METHOD0(TimeScale, uint32_t());
+};
+
+}  // namespace
+
+class TsSegmenterTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    mock_ts_writer_.reset(new MockTsWriter());
+    mock_pes_packet_generator_.reset(new MockPesPacketGenerator());
+  }
+
+  scoped_ptr<MockTsWriter> mock_ts_writer_;
+  scoped_ptr<MockPesPacketGenerator> mock_pes_packet_generator_;
+};
+
+TEST_F(TsSegmenterTest, Initialize) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+}
+
+TEST_F(TsSegmenterTest, AddSample) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_duration = 10.0;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  ON_CALL(*mock_ts_writer_, TimeScale())
+      .WillByDefault(Return(kTimeScale));
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  const uint8_t kAnyData[] = {
+      0x01, 0x0F, 0x3C,
+  };
+  scoped_refptr<MediaSample> sample =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+
+  Sequence writer_sequence;
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts")))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_))
+      .WillOnce(Return(true));
+
+  Sequence ready_pes_sequence;
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+
+  EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
+      .WillOnce(Return(true));
+
+  // The pointer is released inside the segmenter.
+  PesPacket* pes = new PesPacket();
+  pes->set_duration(kTimeScale);
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .WillOnce(Return(pes));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+  EXPECT_OK(segmenter.AddSample(sample));
+}
+
+// Verify the case where the segment is long enough and the current segment
+// should be closed.
+TEST_F(TsSegmenterTest, PastSegmentDuration) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_duration = 10.0;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  ON_CALL(*mock_ts_writer_, TimeScale())
+      .WillByDefault(Return(kTimeScale));
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  const uint8_t kAnyData[] = {
+      0x01, 0x0F, 0x3C,
+  };
+  scoped_refptr<MediaSample> sample =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+
+  Sequence writer_sequence;
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts")))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_)).WillOnce(Return(true));
+
+  Sequence ready_pes_sequence;
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+
+  EXPECT_CALL(*mock_ts_writer_, FinalizeSegment())
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
+      .WillOnce(Return(true));
+
+  // The pointer is released inside the segmenter.
+  PesPacket* pes = new PesPacket();
+  // 11 seconds > 10 seconds (segment duration).
+  // Expect the segment to be finalized.
+  pes->set_duration(11 * kTimeScale);
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .WillOnce(Return(pes));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+  EXPECT_OK(segmenter.AddSample(sample));
+}
+
+// Finalize right after Initialize(). The writer will not be initialized.
+TEST_F(TsSegmenterTest, InitializeThenFinalize) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_duration = 10.0;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, Flush()).WillOnce(Return(true));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+  EXPECT_OK(segmenter.Finalize());
+}
+
+// Verify the "normal" case where samples have been added and the writer has
+// been initialized.
+// The test does not really add any samples but instead simulates an initialized
+// writer with a mock.
+TEST_F(TsSegmenterTest, Finalize) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_duration = 10.0;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  InSequence s;
+  EXPECT_CALL(*mock_pes_packet_generator_, Flush()).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .WillOnce(Return(0u));
+  EXPECT_CALL(*mock_ts_writer_, FinalizeSegment()).WillOnce(Return(true));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+  segmenter.SetTsWriterFileOpenedForTesting(true);
+  EXPECT_OK(segmenter.Finalize());
+}
+
+// Verify that it can generate multiple segments.
+TEST_F(TsSegmenterTest, MultipleSegments) {
+  scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
+      kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
+      kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
+      kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
+  MuxerOptions options;
+  options.segment_duration = 10.0;
+  options.segment_template = "file$Number$.ts";
+  TsSegmenter segmenter(options);
+
+  ON_CALL(*mock_ts_writer_, TimeScale())
+      .WillByDefault(Return(kTimeScale));
+
+  EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
+      .WillOnce(Return(true));
+
+  const uint8_t kAnyData[] = {
+      0x01, 0x0F, 0x3C,
+  };
+  scoped_refptr<MediaSample> sample =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+
+  EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_))
+      .Times(2)
+      .WillRepeatedly(Return(true));
+
+  Sequence writer_sequence;
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts")))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  Sequence ready_pes_sequence;
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+
+  // The pointer is released inside the segmenter.
+  PesPacket* pes = new PesPacket();
+  // 11 seconds > 10 seconds (segment duration).
+  // Expect the segment to be finalized.
+  pes->set_duration(11 * kTimeScale);
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(pes));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+
+  EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_ts_writer_, FinalizeSegment())
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  // Expectations for second AddSample() for the second segment.
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file2.ts")))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+
+  // The pointer is released inside the segmenter.
+  pes = new PesPacket();
+  // 7 < 10 seconds, If FinalizeSegment() is called AddSample will fail (due to
+  // mock returning false by default).
+  pes->set_duration(7 * kTimeScale);
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(pes));
+
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+
+  EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
+
+  segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
+  segmenter.InjectPesPacketGeneratorForTesting(
+      mock_pes_packet_generator_.Pass());
+  EXPECT_OK(segmenter.Initialize(*stream_info));
+  EXPECT_OK(segmenter.AddSample(sample));
+  EXPECT_OK(segmenter.AddSample(sample));
+}
+
+}  // namespace mp2t
+}  // namespace media
+}  // namespace edash_packager

packager/media/formats/mp2t/ts_writer.h

@@ -41,27 +41,27 @@ class ContinuityCounter {
 class TsWriter {
  public:
   TsWriter();
-  ~TsWriter();
+  virtual ~TsWriter();
 
   /// This must be called before calling other methods.
   /// @return true on success, false otherwise.
-  bool Initialize(const StreamInfo& stream_info);
+  virtual bool Initialize(const StreamInfo& stream_info);
 
   /// This will fail if the current segment is not finalized.
   /// @param file_name is the output file name.
   /// @return true on success, false otherwise.
-  bool NewSegment(const std::string& file_name);
+  virtual bool NewSegment(const std::string& file_name);
 
   /// Flush all the pending PesPackets that have not been written to file and
   /// close the file.
   /// @return true on success, false otherwise.
-  bool FinalizeSegment();
+  virtual bool FinalizeSegment();
 
   /// Add PesPacket to the instance. PesPacket might not get written to file
   /// immediately.
   /// @param pes_packet gets added to the writer.
   /// @return true on success, false otherwise.
-  bool AddPesPacket(scoped_ptr<PesPacket> pes_packet);
+  virtual bool AddPesPacket(scoped_ptr<PesPacket> pes_packet);
 
  private:
   std::vector<uint8_t> psi_ts_packets_;