diff --git a/test/test_helper.h b/test/test_helper.h
index 2bd4cf0..599ce5d 100644
--- a/test/test_helper.h
+++ b/test/test_helper.h
@@ -43,11 +43,13 @@
   FakeProducer* ConnectFakeProducer();
   void ConnectConsumer();
   void StartTracing(const TraceConfig& config);
+  void DisableTracing();
+  void FlushAndWait(uint32_t timeout_ms);
   void ReadData(uint32_t read_count = 0);
 
   void WaitForConsumerConnect();
   void WaitForProducerEnabled();
-  void WaitForTracingDisabled();
+  void WaitForTracingDisabled(uint32_t timeout_ms = 5000);
   void WaitForReadData(uint32_t read_count = 0);
 
   std::function<void()> WrapTask(const std::function<void()>& function);