Fix packetization logic to leave space for extensions in the last packet

webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc

Index: webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc
index c53a2ca5a2f112097060eca3b79af1077ecc4ea6..0b316162c025567ee9079b87d7cb835589b43b6f 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_format_vp9_unittest.cc
@@ -138,14 +138,16 @@ class RtpPacketizerVp9Test : public ::testing::Test {
   size_t payload_pos_;
   RTPVideoHeaderVP9 expected_;
   std::unique_ptr<RtpPacketizerVp9> packetizer_;
+  size_t num_packets_;
 
   void Init(size_t payload_size, size_t packet_size) {
     payload_.reset(new uint8_t[payload_size]);
     memset(payload_.get(), 7, payload_size);
     payload_size_ = payload_size;
     payload_pos_ = 0;
-    packetizer_.reset(new RtpPacketizerVp9(expected_, packet_size));
-    packetizer_->SetPayloadData(payload_.get(), payload_size_, NULL);
+    packetizer_.reset(new RtpPacketizerVp9(expected_, packet_size, 0));
+    num_packets_ =
+        packetizer_->SetPayloadData(payload_.get(), payload_size_, NULL);
   }
 
   void CheckPayload(const uint8_t* packet,
@@ -162,24 +164,23 @@ class RtpPacketizerVp9Test : public ::testing::Test {
                                   const size_t* expected_sizes,
                                   size_t expected_num_packets) {
     ASSERT_TRUE(packetizer_.get() != NULL);
-    bool last = false;
     if (expected_num_packets == 0) {
-      EXPECT_FALSE(packetizer_->NextPacket(&packet_, &last));
+      EXPECT_FALSE(packetizer_->NextPacket(&packet_));
       return;
     }
+    EXPECT_EQ(expected_num_packets, num_packets_);
     for (size_t i = 0; i < expected_num_packets; ++i) {
-      EXPECT_TRUE(packetizer_->NextPacket(&packet_, &last));
+      EXPECT_TRUE(packetizer_->NextPacket(&packet_));
       auto rtp_payload = packet_.payload();
       EXPECT_EQ(expected_sizes[i], rtp_payload.size());
       RTPVideoHeaderVP9 hdr = expected_;
       hdr.beginning_of_frame = (i == 0);
-      hdr.end_of_frame = last;
+      hdr.end_of_frame = (i + 1) == expected_num_packets;
       ParseAndCheckPacket(rtp_payload.data(), hdr, expected_hdr_sizes[i],
                           rtp_payload.size());
       CheckPayload(rtp_payload.data(), expected_hdr_sizes[i],
-                   rtp_payload.size(), last);
+                   rtp_payload.size(), (i + 1) == expected_num_packets);
     }
-    EXPECT_TRUE(last);
   }
 };
 
@@ -444,36 +445,32 @@ TEST_F(RtpPacketizerVp9Test, TestOnlyHighestSpatialLayerSetMarker) {
   vp9_header.num_spatial_layers = 3;
 
   RtpPacketToSend packet(kNoExtensions);
-  bool last;
 
   vp9_header.spatial_idx = 0;
-  RtpPacketizerVp9 packetizer0(vp9_header, kPacketSize);
-  packetizer0.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation);
-  ASSERT_TRUE(packetizer0.NextPacket(&packet, &last));
-  EXPECT_FALSE(last);
+  RtpPacketizerVp9 packetizer0(vp9_header, kPacketSize, 0);
+  EXPECT_EQ(
+      packetizer0.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation), 2u);
+  ASSERT_TRUE(packetizer0.NextPacket(&packet));
   EXPECT_FALSE(packet.Marker());
-  ASSERT_TRUE(packetizer0.NextPacket(&packet, &last));
-  EXPECT_TRUE(last);
+  ASSERT_TRUE(packetizer0.NextPacket(&packet));
   EXPECT_FALSE(packet.Marker());
 
   vp9_header.spatial_idx = 1;
-  RtpPacketizerVp9 packetizer1(vp9_header, kPacketSize);
-  packetizer1.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation);
-  ASSERT_TRUE(packetizer1.NextPacket(&packet, &last));
-  EXPECT_FALSE(last);
+  RtpPacketizerVp9 packetizer1(vp9_header, kPacketSize, 0);
+  EXPECT_EQ(
+      packetizer1.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation), 2u);
+  ASSERT_TRUE(packetizer1.NextPacket(&packet));
   EXPECT_FALSE(packet.Marker());
-  ASSERT_TRUE(packetizer1.NextPacket(&packet, &last));
-  EXPECT_TRUE(last);
+  ASSERT_TRUE(packetizer1.NextPacket(&packet));
   EXPECT_FALSE(packet.Marker());
 
   vp9_header.spatial_idx = 2;
-  RtpPacketizerVp9 packetizer2(vp9_header, kPacketSize);
-  packetizer2.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation);
-  ASSERT_TRUE(packetizer2.NextPacket(&packet, &last));
-  EXPECT_FALSE(last);
+  RtpPacketizerVp9 packetizer2(vp9_header, kPacketSize, 0);
+  EXPECT_EQ(
+      packetizer2.SetPayloadData(kFrame, sizeof(kFrame), kNoFragmentation), 2u);
+  ASSERT_TRUE(packetizer2.NextPacket(&packet));
   EXPECT_FALSE(packet.Marker());
-  ASSERT_TRUE(packetizer2.NextPacket(&packet, &last));
-  EXPECT_TRUE(last);
+  ASSERT_TRUE(packetizer2.NextPacket(&packet));
   EXPECT_TRUE(packet.Marker());
 }