From 3d5ea505199003d20ecfb82df898ebf921d4f6c7 Mon Sep 17 00:00:00 2001
From: Grzegorz Latosinski <glatosinski@antmicro.com>
Date: Thu, 7 Dec 2023 16:16:59 +0100
Subject: [PATCH] [#52286] frontend_tester: Added possibility to include header
 to terminal stress test

Signed-off-by: Grzegorz Latosinski <glatosinski@antmicro.com>
---
 .../frontend_tester/frontend_tester_specification.json    | 5 +++++
 pipeline_manager/frontend_tester/tester_client.py         | 8 +++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/pipeline_manager/frontend_tester/frontend_tester_specification.json b/pipeline_manager/frontend_tester/frontend_tester_specification.json
index 86e0b5b9..4847449f 100644
--- a/pipeline_manager/frontend_tester/frontend_tester_specification.json
+++ b/pipeline_manager/frontend_tester/frontend_tester_specification.json
@@ -15,6 +15,11 @@
                     "type": "integer",
                     "default": 10
                 },
+                {
+                    "name": "AddMessageID",
+                    "type": "bool",
+                    "default": false
+                },
                 {
                     "name": "MessagesPerSecond",
                     "type": "number",
diff --git a/pipeline_manager/frontend_tester/tester_client.py b/pipeline_manager/frontend_tester/tester_client.py
index 171f9b7a..cb394724 100644
--- a/pipeline_manager/frontend_tester/tester_client.py
+++ b/pipeline_manager/frontend_tester/tester_client.py
@@ -38,6 +38,7 @@
 from typing import Dict, List, Union
 import string
 import sys
+from datetime import datetime
 
 from pipeline_manager import frontend_tester
 from pipeline_manager.utils.logger import string_to_verbosity
@@ -385,6 +386,7 @@ async def _run_validate_response(
             terminal_name = properties["TerminalName"]
             message_length = properties["MessageLength"]
             rate = properties["MessagesPerSecond"]
+            add_message_id = properties["AddMessageID"]
 
             await self.client.notify(
                 'progress_change',
@@ -397,11 +399,15 @@ async def _run_validate_response(
                     list(string.ascii_uppercase + string.digits) + ['\r\n'],
                     k=message_length
                 ))
+                if add_message_id:
+                    curr_time = datetime.now().strftime("%H:%M:%S.%f")
+                    random_message = f"\r\nMessage {counter} [{curr_time}] :  {random_message}"  # noqa: E501
                 await self.client.notify(
                     'terminal_write',
                     {
                         'name': terminal_name,
-                        'message': f'{random_message}'}
+                        'message': random_message
+                    }
                 )
                 counter += 1
                 await asyncio.sleep(1 / rate)