perfetto: further cleanup and consolidate test code

With the introduction of test helper, integration tests have already been made much cleaner.
Now the next step is to take this further and reduce coupling between producer
service and consumer. We can do this by making the producer emit data without
the test explicitly requesting it (like ftrace) if a flag in config is enabled.

Moreover, we can also consolidate all checkpointing code inside the helper class instead
of every file defining its own.

Bug: 74380167
Change-Id: Idbe24586a4aaadb17f934f6e06c5c6d378be0ec3

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_;