From 5b1bd4b0886a230acc2639ee7ae6a911bc60e655 Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Wed, 25 Sep 2024 18:04:23 +0000
Subject: [PATCH] refactor: dont mix aub and hw wait prints

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 .../command_stream_receiver.cpp               |  9 +++++++
 .../command_stream_receiver_tests.cpp         | 26 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/shared/source/command_stream/command_stream_receiver.cpp b/shared/source/command_stream/command_stream_receiver.cpp
index af3593515b..2728434e89 100644
--- a/shared/source/command_stream/command_stream_receiver.cpp
+++ b/shared/source/command_stream/command_stream_receiver.cpp
@@ -1090,6 +1090,14 @@ bool CommandStreamReceiver::createPerDssBackedBuffer(Device &device) {
 }
 
 void CommandStreamReceiver::printTagAddressContent(TaskCountType taskCountToWait, int64_t waitTimeout, bool start) {
+    if (getType() == NEO::CommandStreamReceiverType::aub) {
+        if (start) {
+            PRINT_DEBUG_STRING(true, stdout, "\nAub dump wait for task count %llu", taskCountToWait);
+        } else {
+            PRINT_DEBUG_STRING(true, stdout, "\nAub dump wait completed.");
+        }
+        return;
+    }
     auto postSyncAddress = getTagAddress();
     if (start) {
         PRINT_DEBUG_STRING(true, stdout,
@@ -1099,6 +1107,7 @@ void CommandStreamReceiver::printTagAddressContent(TaskCountType taskCountToWait
         PRINT_DEBUG_STRING(true, stdout,
                            "%s", "\nWaiting completed. Current value:");
     }
+
     for (uint32_t i = 0; i < activePartitions; i++) {
         PRINT_DEBUG_STRING(true, stdout, " %u", *postSyncAddress);
         postSyncAddress = ptrOffset(postSyncAddress, this->immWritePostSyncWriteOffset);
diff --git a/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp b/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
index 240ddca969..81d6493a0a 100644
--- a/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
+++ b/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
@@ -2519,6 +2519,32 @@ HWTEST_F(CommandStreamReceiverTest, givenMultipleActivePartitionsWhenWaitLogIsEn
     EXPECT_STREQ(expectedOutput.str().c_str(), output.c_str());
 }
 
+HWTEST_F(CommandStreamReceiverTest, givenAubCsrWhenLogWaitingForCompletionEnabledThenSkipFullPrint) {
+    DebugManagerStateRestore restorer;
+    debugManager.flags.LogWaitingForCompletion.set(true);
+
+    auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
+    csr.commandStreamReceiverType = NEO::CommandStreamReceiverType::aub;
+
+    WaitParams waitParams;
+    waitParams.waitTimeout = std::numeric_limits<int64_t>::max();
+
+    testing::internal::CaptureStdout();
+
+    WaitStatus status = csr.waitForCompletionWithTimeout(waitParams, 0);
+    EXPECT_EQ(WaitStatus::ready, status);
+
+    std::string output = testing::internal::GetCapturedStdout();
+
+    std::string expectedOutput1 = "Aub dump wait for task count";
+    std::string expectedOutput2 = "Aub dump wait completed";
+    std::string notExpectedOutput = "Waiting for task count";
+
+    EXPECT_NE(std::string::npos, output.find(expectedOutput1));
+    EXPECT_NE(std::string::npos, output.find(expectedOutput2));
+    EXPECT_EQ(std::string::npos, output.find(notExpectedOutput));
+}
+
 TEST_F(CommandStreamReceiverTest, givenPreambleFlagIsSetWhenGettingFlagStateThenExpectCorrectState) {
     EXPECT_FALSE(commandStreamReceiver->getPreambleSetFlag());
     commandStreamReceiver->setPreambleSetFlag(true);