Merge "perfetto: further cleanup and consolidate test code" am: 563f7f6653

am: 9b180bff85

Change-Id: I43eee0b2e831182a562ddf329a551649da3fb4d7

12 files changed, 146 insertions, 123 deletions

diff --git a/include/perfetto/tracing/core/test_config.h b/include/perfetto/tracing/core/test_config.h
index 0011509fb..c009bd977 100644
--- a/include/perfetto/tracing/core/test_config.h
+++ b/include/perfetto/tracing/core/test_config.h
@@ -68,14 +68,20 @@ class PERFETTO_EXPORT TestConfig {
   uint32_t seed() const { return seed_; }
   void set_seed(uint32_t value) { seed_ = value; }
 
-  uint32_t message_size() const { return message_size_; }
-  void set_message_size(uint32_t value) { message_size_ = value; }
+  uint64_t message_size() const { return message_size_; }
+  void set_message_size(uint64_t value) { message_size_ = value; }
+
+  bool send_batch_on_register() const { return send_batch_on_register_; }
+  void set_send_batch_on_register(bool value) {
+    send_batch_on_register_ = value;
+  }
 
  private:
   uint32_t message_count_ = {};
   uint32_t max_messages_per_second_ = {};
   uint32_t seed_ = {};
-  uint32_t message_size_ = {};
+  uint64_t message_size_ = {};
+  bool send_batch_on_register_ = {};
 
   // Allows to preserve unknown protobuf fields for compatibility
   // with future versions of .proto files.
diff --git a/protos/perfetto/config/perfetto_config.proto b/protos/perfetto/config/perfetto_config.proto
index b2e7ec187..5cddc6216 100644
--- a/protos/perfetto/config/perfetto_config.proto
+++ b/protos/perfetto/config/perfetto_config.proto
@@ -160,7 +160,11 @@ message TestConfig {
   // The size of each message in bytes. Should be greater than or equal 5 to
   // account for the number of bytes needed to encode the random number and a
   // null byte for the string.
-  optional uint32 message_size = 4;
+  optional uint64 message_size = 4;
+
+  // Whether the producer should send a event batch when the data source is
+  // is initially registered.
+  optional bool send_batch_on_register = 5;
 }
 
 // End of protos/perfetto/config/test_config.proto
diff --git a/protos/perfetto/config/test_config.proto b/protos/perfetto/config/test_config.proto
index 52bd9ef46..2c205de91 100644
--- a/protos/perfetto/config/test_config.proto
+++ b/protos/perfetto/config/test_config.proto
@@ -40,5 +40,9 @@ message TestConfig {
   // The size of each message in bytes. Should be greater than or equal 5 to
   // account for the number of bytes needed to encode the random number and a
   // null byte for the string.
-  optional uint32 message_size = 4;
+  optional uint64 message_size = 4;
+
+  // Whether the producer should send a event batch when the data source is
+  // is initially registered.
+  optional bool send_batch_on_register = 5;
 }
diff --git a/src/tracing/core/test_config.cc b/src/tracing/core/test_config.cc
index d58cbcee3..6aa26ab23 100644
--- a/src/tracing/core/test_config.cc
+++ b/src/tracing/core/test_config.cc
@@ -55,6 +55,12 @@ void TestConfig::FromProto(const perfetto::protos::TestConfig& proto) {
   static_assert(sizeof(message_size_) == sizeof(proto.message_size()),
                 "size mismatch");
   message_size_ = static_cast<decltype(message_size_)>(proto.message_size());
+
+  static_assert(
+      sizeof(send_batch_on_register_) == sizeof(proto.send_batch_on_register()),
+      "size mismatch");
+  send_batch_on_register_ = static_cast<decltype(send_batch_on_register_)>(
+      proto.send_batch_on_register());
   unknown_fields_ = proto.unknown_fields();
 }
 
@@ -80,6 +86,13 @@ void TestConfig::ToProto(perfetto::protos::TestConfig* proto) const {
                 "size mismatch");
   proto->set_message_size(
       static_cast<decltype(proto->message_size())>(message_size_));
+
+  static_assert(sizeof(send_batch_on_register_) ==
+                    sizeof(proto->send_batch_on_register()),
+                "size mismatch");
+  proto->set_send_batch_on_register(
+      static_cast<decltype(proto->send_batch_on_register())>(
+          send_batch_on_register_));
   *(proto->mutable_unknown_fields()) = unknown_fields_;
 }
 
diff --git a/test/cts/end_to_end_integrationtest_cts.cc b/test/cts/end_to_end_integrationtest_cts.cc
index 0c47c10e9..439227189 100644
--- a/test/cts/end_to_end_integrationtest_cts.cc
+++ b/test/cts/end_to_end_integrationtest_cts.cc
@@ -36,6 +36,7 @@ class PerfettoCtsTest : public ::testing::Test {
 
     TestHelper helper(&task_runner);
     helper.ConnectConsumer();
+    helper.WaitForConsumerConnect();
 
     TraceConfig trace_config;
     trace_config.add_buffers()->set_size_kb(1024);
@@ -51,28 +52,22 @@ class PerfettoCtsTest : public ::testing::Test {
     ds_config->mutable_for_testing()->set_seed(kRandomSeed);
     ds_config->mutable_for_testing()->set_message_count(kEventCount);
     ds_config->mutable_for_testing()->set_message_size(kMessageSizeBytes);
+    ds_config->mutable_for_testing()->set_send_batch_on_register(true);
 
-    auto producer_enabled = task_runner.CreateCheckpoint("producer.enabled");
-    task_runner.PostTask(producer_enabled);
     helper.StartTracing(trace_config);
+    helper.WaitForTracingDisabled();
+
+    helper.ReadData();
+    helper.WaitForReadData();
+
+    const auto& packets = helper.trace();
+    ASSERT_EQ(packets.size(), kEventCount);
 
-    size_t packets_seen = 0;
     std::minstd_rand0 rnd_engine(kRandomSeed);
-    auto on_consumer_data = [&packets_seen,
-                             &rnd_engine](const protos::TracePacket& packet) {
+    for (const auto& packet : packets) {
       ASSERT_TRUE(packet.has_for_testing());
       ASSERT_EQ(packet.for_testing().seq_value(), rnd_engine());
-      packets_seen++;
-    };
-    auto on_readback_complete =
-        task_runner.CreateCheckpoint("readback.complete");
-    task_runner.PostDelayedTask(
-        [&on_consumer_data, &on_readback_complete, &helper]() {
-          helper.ReadData(on_consumer_data, on_readback_complete);
-        },
-        1000);
-    task_runner.RunUntilCheckpoint("readback.complete");
-    ASSERT_EQ(packets_seen, kEventCount);
+    }
   }
 };
 
diff --git a/test/end_to_end_benchmark.cc b/test/end_to_end_benchmark.cc
index 607c73559..100674e98 100644
--- a/test/end_to_end_benchmark.cc
+++ b/test/end_to_end_benchmark.cc
@@ -44,17 +44,15 @@ void BenchmarkCommon(benchmark::State& state) {
 
   FakeProducer* producer = helper.ConnectFakeProducer();
   helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
 
-  // Setup the TraceConfig for the consumer.
   TraceConfig trace_config;
   trace_config.add_buffers()->set_size_kb(512);
 
-  // Create the buffer for ftrace.
   auto* ds_config = trace_config.add_data_sources()->mutable_config();
   ds_config->set_name("android.perfetto.FakeProducer");
   ds_config->set_target_buffer(0);
 
-  // The parameters for the producer.
   static constexpr uint32_t kRandomSeed = 42;
   uint32_t message_count = static_cast<uint32_t>(state.range(0));
   uint32_t message_bytes = static_cast<uint32_t>(state.range(1));
@@ -64,26 +62,13 @@ void BenchmarkCommon(benchmark::State& state) {
   uint32_t time_for_messages_ms =
       10000 + (messages_per_s == 0 ? 0 : message_count * 1000 / messages_per_s);
 
-  // Setup the test to use a random number generator.
   ds_config->mutable_for_testing()->set_seed(kRandomSeed);
   ds_config->mutable_for_testing()->set_message_count(message_count);
   ds_config->mutable_for_testing()->set_message_size(message_bytes);
   ds_config->mutable_for_testing()->set_max_messages_per_second(messages_per_s);
 
   helper.StartTracing(trace_config);
-
-  bool is_first_packet = true;
-  std::minstd_rand0 rnd_engine(kRandomSeed);
-  auto on_consumer_data = [&is_first_packet,
-                           &rnd_engine](const protos::TracePacket& packet) {
-    ASSERT_TRUE(packet.has_for_testing());
-    if (is_first_packet) {
-      rnd_engine = std::minstd_rand0(packet.for_testing().seq_value());
-      is_first_packet = false;
-    } else {
-      ASSERT_EQ(packet.for_testing().seq_value(), rnd_engine());
-    }
-  };
+  helper.WaitForProducerEnabled();
 
   uint64_t wall_start_ns = static_cast<uint64_t>(base::GetWallTimeNs().count());
   uint64_t service_start_ns = helper.service_thread()->GetThreadCPUTimeNs();
@@ -106,12 +91,23 @@ void BenchmarkCommon(benchmark::State& state) {
   state.counters["Ser CPU"] = benchmark::Counter(100.0 * service_ns / wall_ns);
   state.counters["Ser ns/m"] =
       benchmark::Counter(1.0 * service_ns / message_count);
+  state.SetBytesProcessed(iterations * message_bytes * message_count);
 
   // Read back the buffer just to check correctness.
-  auto on_readback_complete = task_runner.CreateCheckpoint("readback.complete");
-  helper.ReadData(on_consumer_data, on_readback_complete);
-  task_runner.RunUntilCheckpoint("readback.complete");
-  state.SetBytesProcessed(iterations * message_bytes * message_count);
+  helper.ReadData();
+  helper.WaitForReadData();
+
+  bool is_first_packet = true;
+  std::minstd_rand0 rnd_engine(kRandomSeed);
+  for (const auto& packet : helper.trace()) {
+    ASSERT_TRUE(packet.has_for_testing());
+    if (is_first_packet) {
+      rnd_engine = std::minstd_rand0(packet.for_testing().seq_value());
+      is_first_packet = false;
+    } else {
+      ASSERT_EQ(packet.for_testing().seq_value(), rnd_engine());
+    }
+  }
 }
 
 void SaturateCpuArgs(benchmark::internal::Benchmark* b) {
diff --git a/test/end_to_end_integrationtest.cc b/test/end_to_end_integrationtest.cc
index 2a8e1f067..289f65dd0 100644
--- a/test/end_to_end_integrationtest.cc
+++ b/test/end_to_end_integrationtest.cc
@@ -67,6 +67,7 @@ TEST(PerfettoTest, MAYBE_TestFtraceProducer) {
 #endif
 
   helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
 
   TraceConfig trace_config;
   trace_config.add_buffers()->set_size_kb(1024);
@@ -80,25 +81,20 @@ TEST(PerfettoTest, MAYBE_TestFtraceProducer) {
   *ftrace_config->add_ftrace_events() = "sched_switch";
   *ftrace_config->add_ftrace_events() = "bar";
 
-  auto producer_enabled = task_runner.CreateCheckpoint("producer.enabled");
-  task_runner.PostDelayedTask(producer_enabled, 100);
   helper.StartTracing(trace_config);
+  helper.WaitForTracingDisabled();
 
-  size_t packets_seen = 0;
-  auto on_consumer_data = [&packets_seen](const protos::TracePacket& packet) {
+  helper.ReadData();
+  helper.WaitForReadData();
+
+  const auto& packets = helper.trace();
+  ASSERT_GT(packets.size(), 0u);
+
+  for (const auto& packet : packets) {
     for (int ev = 0; ev < packet.ftrace_events().event_size(); ev++) {
       ASSERT_TRUE(packet.ftrace_events().event(ev).has_sched_switch());
     }
-    packets_seen++;
-  };
-  auto on_readback_complete = task_runner.CreateCheckpoint("readback.complete");
-  task_runner.PostDelayedTask(
-      [&helper, &on_consumer_data, &on_readback_complete] {
-        helper.ReadData(on_consumer_data, on_readback_complete);
-      },
-      3000);
-  task_runner.RunUntilCheckpoint("readback.complete");
-  ASSERT_GT(packets_seen, 0u);
+  }
 }
 
 TEST(PerfettoTest, TestFakeProducer) {
@@ -106,9 +102,9 @@ TEST(PerfettoTest, TestFakeProducer) {
 
   TestHelper helper(&task_runner);
   helper.StartServiceIfRequired();
-
-  FakeProducer* producer = helper.ConnectFakeProducer();
+  helper.ConnectFakeProducer();
   helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
 
   TraceConfig trace_config;
   trace_config.add_buffers()->set_size_kb(1024);
@@ -124,25 +120,22 @@ TEST(PerfettoTest, TestFakeProducer) {
   ds_config->mutable_for_testing()->set_seed(kRandomSeed);
   ds_config->mutable_for_testing()->set_message_count(kNumPackets);
   ds_config->mutable_for_testing()->set_message_size(kMsgSize);
+  ds_config->mutable_for_testing()->set_send_batch_on_register(true);
 
   helper.StartTracing(trace_config);
+  helper.WaitForTracingDisabled();
 
-  producer->ProduceEventBatch(
-      helper.WrapTask(task_runner.CreateCheckpoint("produced.and.committed")));
-  task_runner.RunUntilCheckpoint("produced.and.committed");
+  helper.ReadData();
+  helper.WaitForReadData();
+
+  const auto& packets = helper.trace();
+  ASSERT_EQ(packets.size(), kNumPackets);
 
-  size_t packets_seen = 0;
   std::minstd_rand0 rnd_engine(kRandomSeed);
-  auto on_consumer_data = [&packets_seen,
-                           &rnd_engine](const protos::TracePacket& packet) {
+  for (const auto& packet : packets) {
     ASSERT_TRUE(packet.has_for_testing());
     ASSERT_EQ(packet.for_testing().seq_value(), rnd_engine());
-    packets_seen++;
-  };
-  auto on_readback_complete = task_runner.CreateCheckpoint("readback.complete");
-  helper.ReadData(on_consumer_data, on_readback_complete);
-  task_runner.RunUntilCheckpoint("readback.complete");
-  ASSERT_EQ(packets_seen, kNumPackets);
+  }
 }
 
 TEST(PerfettoTest, VeryLargePackets) {
@@ -150,13 +143,13 @@ TEST(PerfettoTest, VeryLargePackets) {
 
   TestHelper helper(&task_runner);
   helper.StartServiceIfRequired();
-
-  FakeProducer* producer = helper.ConnectFakeProducer();
+  helper.ConnectFakeProducer();
   helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
 
-  // Setup the TraceConfig for the consumer.
   TraceConfig trace_config;
   trace_config.add_buffers()->set_size_kb(4096 * 10);
+  trace_config.set_duration_ms(500);
 
   auto* ds_config = trace_config.add_data_sources()->mutable_config();
   ds_config->set_name("android.perfetto.FakeProducer");
@@ -168,29 +161,26 @@ TEST(PerfettoTest, VeryLargePackets) {
   ds_config->mutable_for_testing()->set_seed(kRandomSeed);
   ds_config->mutable_for_testing()->set_message_count(kNumPackets);
   ds_config->mutable_for_testing()->set_message_size(kMsgSize);
+  ds_config->mutable_for_testing()->set_send_batch_on_register(true);
 
   helper.StartTracing(trace_config);
+  helper.WaitForTracingDisabled();
+
+  helper.ReadData();
+  helper.WaitForReadData();
 
-  producer->ProduceEventBatch(
-      helper.WrapTask(task_runner.CreateCheckpoint("produced.and.committed")));
-  task_runner.RunUntilCheckpoint("produced.and.committed");
+  const auto& packets = helper.trace();
+  ASSERT_EQ(packets.size(), kNumPackets);
 
-  size_t packets_seen = 0;
   std::minstd_rand0 rnd_engine(kRandomSeed);
-  auto on_consumer_data = [&packets_seen,
-                           &rnd_engine](const protos::TracePacket& packet) {
+  for (const auto& packet : packets) {
     ASSERT_TRUE(packet.has_for_testing());
     ASSERT_EQ(packet.for_testing().seq_value(), rnd_engine());
     size_t msg_size = packet.for_testing().str().size();
     ASSERT_EQ(kMsgSize, msg_size);
     for (size_t i = 0; i < msg_size; i++)
       ASSERT_EQ(i < msg_size - 1 ? '.' : 0, packet.for_testing().str()[i]);
-    packets_seen++;
-  };
-  auto on_readback_complete = task_runner.CreateCheckpoint("readback.complete");
-  helper.ReadData(on_consumer_data, on_readback_complete);
-  task_runner.RunUntilCheckpoint("readback.complete");
-  ASSERT_EQ(packets_seen, kNumPackets);
+  }
 }
 
 }  // namespace perfetto
diff --git a/test/end_to_end_shared_memory_fuzzer.cc b/test/end_to_end_shared_memory_fuzzer.cc
index b509647b0..b9151942a 100644
--- a/test/end_to_end_shared_memory_fuzzer.cc
+++ b/test/end_to_end_shared_memory_fuzzer.cc
@@ -139,17 +139,14 @@ int FuzzSharedMemory(const uint8_t* data, size_t size) {
   TestHelper helper(&task_runner);
   helper.StartServiceIfRequired();
 
-  auto on_produced_and_committed =
-      task_runner.CreateCheckpoint("produced.and.committed");
-  auto posted_on_produced_and_committed = [&task_runner,
-                                           &on_produced_and_committed] {
-    task_runner.PostTask(on_produced_and_committed);
-  };
   TaskRunnerThread producer_thread("perfetto.prd");
   producer_thread.Start(std::unique_ptr<FakeProducerDelegate>(
-      new FakeProducerDelegate(data, size, posted_on_produced_and_committed)));
+      new FakeProducerDelegate(data, size,
+                               helper.WrapTask(task_runner.CreateCheckpoint(
+                                   "produced.and.committed")))));
 
   helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
 
   TraceConfig trace_config;
   trace_config.add_buffers()->set_size_kb(8);
@@ -158,19 +155,11 @@ int FuzzSharedMemory(const uint8_t* data, size_t size) {
   ds_config->set_name("android.perfetto.FakeProducer");
   ds_config->set_target_buffer(0);
 
-  auto producer_enabled = task_runner.CreateCheckpoint("producer.enabled");
-  task_runner.PostTask(producer_enabled);
   helper.StartTracing(trace_config);
   task_runner.RunUntilCheckpoint("produced.and.committed");
 
-  auto on_readback_complete = task_runner.CreateCheckpoint("readback.complete");
-  auto on_consumer_data =
-      [&on_readback_complete](const protos::TracePacket& packet) {
-        if (packet.for_testing().str() == "end")
-          on_readback_complete();
-      };
-  helper.ReadData(on_consumer_data, [] {});
-  task_runner.RunUntilCheckpoint("readback.complete");
+  helper.ReadData();
+  helper.WaitForReadData();
 
   return 0;
 }
diff --git a/test/fake_producer.cc b/test/fake_producer.cc
index 041e2dc56..f2922d630 100644
--- a/test/fake_producer.cc
+++ b/test/fake_producer.cc
@@ -68,7 +68,11 @@ void FakeProducer::CreateDataSourceInstance(
   message_size_ = source_config.for_testing().message_size();
   max_messages_per_second_ =
       source_config.for_testing().max_messages_per_second();
-  task_runner_->PostTask(on_create_data_source_instance_);
+  if (source_config.for_testing().send_batch_on_register()) {
+    ProduceEventBatch(on_create_data_source_instance_);
+  } else {
+    task_runner_->PostTask(on_create_data_source_instance_);
+  }
 }
 
 void FakeProducer::TearDownDataSourceInstance(DataSourceInstanceID) {
@@ -76,7 +80,7 @@ void FakeProducer::TearDownDataSourceInstance(DataSourceInstanceID) {
   trace_writer_.reset();
 }
 
-// Note: this will called on a different thread.
+// Note: this can be called on a different thread.
 void FakeProducer::ProduceEventBatch(std::function<void()> callback) {
   task_runner_->PostTask([this, callback] {
     PERFETTO_CHECK(trace_writer_);
diff --git a/test/fake_producer.h b/test/fake_producer.h
index 35149e3df..622ed94e8 100644
--- a/test/fake_producer.h
+++ b/test/fake_producer.h
@@ -41,7 +41,7 @@ class FakeProducer : public Producer {
 
   // Produces a batch of events (as configured in the DataSourceConfig) and
   // posts a callback when the service acknowledges the commit.
-  void ProduceEventBatch(std::function<void()> callback);
+  void ProduceEventBatch(std::function<void()> callback = [] {});
 
   // Producer implementation.
   void OnConnect() override;
@@ -59,7 +59,7 @@ class FakeProducer : public Producer {
   base::TaskRunner* task_runner_ = nullptr;
   std::string name_;
   std::minstd_rand0 rnd_engine_;
-  uint32_t message_size_ = 0;
+  uint64_t message_size_ = 0;
   uint32_t message_count_ = 0;
   uint32_t max_messages_per_second_ = 0;
   std::function<void()> on_create_data_source_instance_;
diff --git a/test/test_helper.cc b/test/test_helper.cc
index 8055b7b33..e5b78c927 100644
--- a/test/test_helper.cc
+++ b/test/test_helper.cc
@@ -45,14 +45,16 @@ TestHelper::TestHelper(base::TestTaskRunner* task_runner)
       producer_thread_("perfetto.prd") {}
 
 void TestHelper::OnConnect() {
-  std::move(continuation_callack_)();
+  std::move(on_connect_callback_)();
 }
 
 void TestHelper::OnDisconnect() {
   FAIL() << "Consumer unexpectedly disconnected from the service";
 }
 
-void TestHelper::OnTracingDisabled() {}
+void TestHelper::OnTracingDisabled() {
+  std::move(on_stop_tracing_callback_)();
+}
 
 void TestHelper::OnTraceData(std::vector<TracePacket> packets, bool has_more) {
   for (auto& encoded_packet : packets) {
@@ -62,12 +64,11 @@ void TestHelper::OnTraceData(std::vector<TracePacket> packets, bool has_more) {
       continue;
     ASSERT_EQ(protos::TracePacket::kTrustedUid,
               packet.optional_trusted_uid_case());
-    packet_callback_(packet);
+    trace_.push_back(std::move(packet));
   }
 
   if (!has_more) {
-    packet_callback_ = {};
-    std::move(continuation_callack_)();
+    std::move(on_packets_finished_callback_)();
   }
 }
 
@@ -89,25 +90,38 @@ FakeProducer* TestHelper::ConnectFakeProducer() {
 }
 
 void TestHelper::ConnectConsumer() {
-  continuation_callack_ = task_runner_->CreateCheckpoint("consumer.connected");
+  on_connect_callback_ = task_runner_->CreateCheckpoint("consumer.connected");
   endpoint_ =
       ConsumerIPCClient::Connect(TEST_CONSUMER_SOCK_NAME, this, task_runner_);
-  task_runner_->RunUntilCheckpoint("consumer.connected");
 }
 
 void TestHelper::StartTracing(const TraceConfig& config) {
+  on_stop_tracing_callback_ = task_runner_->CreateCheckpoint("stop.tracing");
   endpoint_->EnableTracing(config);
-  task_runner_->RunUntilCheckpoint("producer.enabled");
 }
 
-void TestHelper::ReadData(
-    std::function<void(const protos::TracePacket&)> packet_callback,
-    std::function<void()> on_finish_callback) {
-  packet_callback_ = packet_callback;
-  continuation_callack_ = on_finish_callback;
+void TestHelper::ReadData() {
+  on_packets_finished_callback_ =
+      task_runner_->CreateCheckpoint("readback.complete");
   endpoint_->ReadBuffers();
 }
 
+void TestHelper::WaitForConsumerConnect() {
+  task_runner_->RunUntilCheckpoint("consumer.connected");
+}
+
+void TestHelper::WaitForProducerEnabled() {
+  task_runner_->RunUntilCheckpoint("producer.enabled");
+}
+
+void TestHelper::WaitForTracingDisabled() {
+  task_runner_->RunUntilCheckpoint("stop.tracing");
+}
+
+void TestHelper::WaitForReadData() {
+  task_runner_->RunUntilCheckpoint("readback.complete");
+}
+
 std::function<void()> TestHelper::WrapTask(
     const std::function<void()>& function) {
   return [this, function] { task_runner_->PostTask(function); };
diff --git a/test/test_helper.h b/test/test_helper.h
index 563345fe8..b26a88d3d 100644
--- a/test/test_helper.h
+++ b/test/test_helper.h
@@ -43,19 +43,27 @@ class TestHelper : public Consumer {
   FakeProducer* ConnectFakeProducer();
   void ConnectConsumer();
   void StartTracing(const TraceConfig& config);
-  void ReadData(std::function<void(const protos::TracePacket&)> packet_callback,
-                std::function<void()> on_finish_callback);
+  void ReadData();
+
+  void WaitForConsumerConnect();
+  void WaitForProducerEnabled();
+  void WaitForTracingDisabled();
+  void WaitForReadData();
 
   std::function<void()> WrapTask(const std::function<void()>& function);
 
   TaskRunnerThread* service_thread() { return &service_thread_; }
   TaskRunnerThread* producer_thread() { return &producer_thread_; }
+  const std::vector<protos::TracePacket>& trace() { return trace_; }
 
  private:
   base::TestTaskRunner* task_runner_ = nullptr;
 
-  std::function<void(const protos::TracePacket&)> packet_callback_;
-  std::function<void()> continuation_callack_;
+  std::function<void()> on_connect_callback_;
+  std::function<void()> on_packets_finished_callback_;
+  std::function<void()> on_stop_tracing_callback_;
+
+  std::vector<protos::TracePacket> trace_;
 
   TaskRunnerThread service_thread_;
   TaskRunnerThread producer_thread_;