From 600d62e15a42d3090d5363a6a1dee0833ab59aa3 Mon Sep 17 00:00:00 2001
From: Ivan Golikov <ivan_golikov@epam.com>
Date: Sun, 21 Apr 2024 20:45:53 +0200
Subject: [PATCH] Enriching failure message from
 task_schedule_css_selector_generation with info about elements that task
 should've processed

---
 app/css_locators/__init__.py |  5 +++--
 app/css_locators/tasks.py    |  3 +++
 utils/api_utils.py           | 12 ++++++++++--
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/app/css_locators/__init__.py b/app/css_locators/__init__.py
index 71deb82c..5161af98 100644
--- a/app/css_locators/__init__.py
+++ b/app/css_locators/__init__.py
@@ -1,7 +1,8 @@
 __all__ = [
-    "task_schedule_css_locator_generation",
+    "CSS_SELECTOR_GEN_TASK_PREFIX",
+    "task_schedule_css_selector_generation",
     "inject_css_selector_generator_scripts",
 ]
 
-from .tasks import task_schedule_css_selector_generation
+from .tasks import task_schedule_css_selector_generation, CSS_SELECTOR_GEN_TASK_PREFIX
 from .utils import inject_css_selector_generator_scripts
diff --git a/app/css_locators/tasks.py b/app/css_locators/tasks.py
index f9a07359..aadc7eb5 100644
--- a/app/css_locators/tasks.py
+++ b/app/css_locators/tasks.py
@@ -9,6 +9,9 @@
 logger = logging.getLogger(__name__)
 
 
+CSS_SELECTOR_GEN_TASK_PREFIX = "css-selectors-gen-"
+
+
 @celery_app.task(bind=True)
 def task_schedule_css_selector_generation(
         self, document_path: str, elements_ids: List[str]
diff --git a/utils/api_utils.py b/utils/api_utils.py
index 93d0881f..1368ad6d 100644
--- a/utils/api_utils.py
+++ b/utils/api_utils.py
@@ -11,7 +11,7 @@
 import app.mongodb as mongodb
 from app.celery_app import celery_app
 from app.constants import CeleryStatuses, WebSocketResponseActions
-from app.css_locators import inject_css_selector_generator_scripts
+from app.css_locators import inject_css_selector_generator_scripts, CSS_SELECTOR_GEN_TASK_PREFIX
 from app.logger import logger
 from app.models import LoggingInfoModel, TaskStatusModel, XPathGenerationModel, CSSSelectorGenerationModel
 from app.redis_app import redis_app
@@ -70,6 +70,14 @@ async def wait_until_task_reach_status(
                 task_dict = task.model_dump()
                 # deleting underscores in task_id if any to send to frontend
                 task_dict["id"] = task_dict["id"].strip("_")
+
+                # Informing frontend about failed elements from
+                # task_schedule_css_locator_generation task, because otherwise
+                # this info will be lost. In schedule_multiple_xpath_generations task
+                # the id of an element is the task id, so it doesn't need this logic
+                if task_id.startswith(CSS_SELECTOR_GEN_TASK_PREFIX):
+                    task_dict["elements_ids"] = task_result_obj.kwargs.get("elements_ids")
+
                 error_message = str(task_result_obj.result)
                 response = get_websocket_response(
                     action=WebSocketResponseActions.STATUS_CHANGED,
@@ -344,7 +352,7 @@ async def process_incoming_ws_request(
 
         for start_idx, end_idx in jobs_chunks:
             task_id = convert_task_id_if_exists(
-                f"css-selectors-gen-{uuid.uuid4()}"
+                f"{CSS_SELECTOR_GEN_TASK_PREFIX}{uuid.uuid4()}"
             )
             task_kwargs = {
                 "document_path": str(html_file),