1 files changed, 5 insertions, 11 deletions
diff --git a/src/profiling/test/ProfilingTests.hpp b/src/profiling/test/ProfilingTests.hpp
index 4d2f974344..21c98723be 100644
--- a/src/profiling/test/ProfilingTests.hpp
+++ b/src/profiling/test/ProfilingTests.hpp
@@ -9,14 +9,12 @@
 
 #include <CommandHandlerFunctor.hpp>
 #include <IProfilingConnection.hpp>
-#include <IProfilingConnectionFactory.hpp>
 #include <Logging.hpp>
 #include <ProfilingService.hpp>
 
 #include <boost/test/unit_test.hpp>
 
 #include <chrono>
-#include <iostream>
 #include <thread>
 
 namespace armnn
@@ -137,15 +135,6 @@ class TestFunctorC : public TestFunctorA
     using TestFunctorA::TestFunctorA;
 };
 
-class MockProfilingConnectionFactory : public IProfilingConnectionFactory
-{
-public:
-    IProfilingConnectionPtr GetProfilingConnection(const ExternalProfilingOptions& options) const override
-    {
-        return std::make_unique<MockProfilingConnection>();
-    }
-};
-
 class SwapProfilingConnectionFactoryHelper : public ProfilingService
 {
 public:
@@ -182,6 +171,11 @@ public:
         TransitionToState(ProfilingService::Instance(), newState);
     }
 
+    void WaitForProfilingPacketsSent()
+    {
+        return WaitForPacketSent(ProfilingService::Instance());
+    }
+
 private:
     MockProfilingConnectionFactoryPtr m_MockProfilingConnectionFactory;
     IProfilingConnectionFactory* m_BackupProfilingConnectionFactory;