From dd8460beba37cbe36d0be8c5e5a2c74fc5b5c0e0 Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Mon, 9 Sep 2024 15:31:32 +0000
Subject: [PATCH] refactor: reduce TBX download timeout for unit tests

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 .../source/command_stream/tbx_command_stream_receiver_hw.h   | 3 +++
 .../source/command_stream/tbx_command_stream_receiver_hw.inl | 3 +--
 shared/test/common/mocks/mock_tbx_csr.h                      | 5 +++++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/shared/source/command_stream/tbx_command_stream_receiver_hw.h b/shared/source/command_stream/tbx_command_stream_receiver_hw.h
index 049056842f..1b3320208e 100644
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.h
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.h
@@ -31,6 +31,9 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     uint32_t getMaskAndValueForPollForCompletion() const;
     bool getpollNotEqualValueForPollForCompletion() const;
     void flushSubmissionsAndDownloadAllocations(TaskCountType taskCount, bool skipAllocationsDownload);
+    MOCKABLE_VIRTUAL uint64_t getNonBlockingDownloadTimeoutMs() const {
+        return 2000; // 2s
+    }
 
   public:
     using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::initAdditionalMMIO;
diff --git a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
index fed032000c..7f15b48c72 100644
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
@@ -578,7 +578,6 @@ void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocationTbx(GraphicsAlloca
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocations(bool blockingWait, TaskCountType taskCount) {
     volatile TagAddressType *pollAddress = this->getTagAddress();
-    constexpr uint64_t timeoutMs = 1000 * 2; // 2s
 
     auto waitTaskCount = std::min(taskCount, this->latestFlushedTaskCount.load());
 
@@ -593,7 +592,7 @@ void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocations(bool blockingWai
                 if (!blockingWait) {
                     // Additional delay to reach PC in case of Event wait
                     timeDiff = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - startTime).count();
-                    if (timeDiff > timeoutMs) {
+                    if (timeDiff > getNonBlockingDownloadTimeoutMs()) {
                         return;
                     }
                 }
diff --git a/shared/test/common/mocks/mock_tbx_csr.h b/shared/test/common/mocks/mock_tbx_csr.h
index fcef4d1317..1448990489 100644
--- a/shared/test/common/mocks/mock_tbx_csr.h
+++ b/shared/test/common/mocks/mock_tbx_csr.h
@@ -111,6 +111,11 @@ struct MockTbxCsrRegisterDownloadedAllocations : TbxCommandStreamReceiverHw<GfxF
         obtainUniqueOwnershipCalled++;
         return TbxCommandStreamReceiverHw<GfxFamily>::obtainUniqueOwnership();
     }
+
+    uint64_t getNonBlockingDownloadTimeoutMs() const override {
+        return 1;
+    }
+
     std::set<GraphicsAllocation *> downloadedAllocations;
     bool flushBatchedSubmissionsCalled = false;
     bool flushTagCalled = false;