3 files changed, 23 insertions, 4 deletions
diff --git a/tests/profiling/gatordmock/CommandFileParser.cpp b/tests/profiling/gatordmock/CommandFileParser.cpp
index d08e72cadd..4a8a19b5d2 100644
--- a/tests/profiling/gatordmock/CommandFileParser.cpp
+++ b/tests/profiling/gatordmock/CommandFileParser.cpp
@@ -23,7 +23,7 @@ void CommandFileParser::ParseFile(std::string CommandFile, GatordMockService& mo
 
     std::cout << "Parsing command file: " << CommandFile << std::endl;
 
-    while (std::getline(infile, line))
+    while (mockService.ReceiveThreadRunning() && std::getline(infile, line))
     {
         std::istringstream iss(line);
 
diff --git a/tests/profiling/gatordmock/GatordMockService.cpp b/tests/profiling/gatordmock/GatordMockService.cpp
index 1cdb024273..529ef063dd 100644
--- a/tests/profiling/gatordmock/GatordMockService.cpp
+++ b/tests/profiling/gatordmock/GatordMockService.cpp
@@ -166,7 +166,11 @@ bool GatordMockService::LaunchReceivingThread()
 
 void GatordMockService::WaitForReceivingThread()
 {
-    m_CloseReceivingThread.store(true);
+    // The receiving thread may already have died.
+    if (m_CloseReceivingThread != true)
+    {
+        m_CloseReceivingThread.store(true);
+    }
     // Check that the receiving thread is running
     if (m_ListeningThread.joinable())
     {
@@ -210,8 +214,16 @@ void GatordMockService::SendPeriodicCounterSelectionList(uint32_t period, std::v
 
 void GatordMockService::WaitCommand(uint timeout)
 {
-    std::this_thread::sleep_for(std::chrono::microseconds(timeout));
-
+    // Wait for a maximum of timeout microseconds or if the receive thread has closed.
+    // There is a certain level of rounding involved in this timing.
+    uint iterations = timeout / 1000;
+    std::cout << std::dec << "Wait command with timeout of " << timeout << " iterations =  " << iterations << std::endl;
+    uint count = 0;
+    while ((this->ReceiveThreadRunning() && (count < iterations)))
+    {
+        std::this_thread::sleep_for(std::chrono::microseconds(1000));
+        ++count;
+    }
     if (m_EchoPackets)
     {
         std::cout << std::dec << "Wait command with timeout of " << timeout << " microseconds completed. " << std::endl;
diff --git a/tests/profiling/gatordmock/GatordMockService.hpp b/tests/profiling/gatordmock/GatordMockService.hpp
index a77d7ff480..c3afc333ca 100644
--- a/tests/profiling/gatordmock/GatordMockService.hpp
+++ b/tests/profiling/gatordmock/GatordMockService.hpp
@@ -41,6 +41,7 @@ public:
     GatordMockService(armnn::profiling::CommandHandlerRegistry& registry, bool echoPackets)
         : m_HandlerRegistry(registry)
         , m_EchoPackets(echoPackets)
+        , m_CloseReceivingThread(false)
     {
         m_PacketsReceivedCount.store(0, std::memory_order_relaxed);
     }
@@ -86,6 +87,12 @@ public:
     /// command handling code is added.
     void WaitForReceivingThread();
 
+    // @return true only if the receive thread is closed or closing.
+    bool ReceiveThreadRunning()
+    {
+        return !m_CloseReceivingThread.load();
+    }
+
     /// Send the counter list to ArmNN.
     void SendPeriodicCounterSelectionList(uint32_t period, std::vector<uint16_t> counters);