From 51beda6bfb12b04d3b0bca87c2cce1f33dd05d0a Mon Sep 17 00:00:00 2001
From: Shivanand Mishra <xemishra@gmail.com>
Date: Sat, 6 Jun 2026 18:25:52 +0530
Subject: [PATCH] fix: defer MockUrlInvoker import and correct invoke() retry
 logic

---
 estimators/factory.py | 12 ++++++++++--
 util/connectors.py    | 17 +++++++++++++----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/estimators/factory.py b/estimators/factory.py
index 9ddc8c61..19464518 100644
--- a/estimators/factory.py
+++ b/estimators/factory.py
@@ -3,7 +3,6 @@
 from estimators.eo_in_place_archive_estimator import EOInPlaceArchiveEstimator
 from estimators.eo_shared_mailbox_estimator import EOSharedMailBoxEstimator
 from estimators.file_estimator import FileEstimator
-from tests.files.mocks import MockUrlInvoker
 from util.auth_manager import TokenManager
 from util.connectors import UrlInvoker
 from util.utils import ScanConfig
@@ -65,6 +64,15 @@ def get_url_invoker(self, hard_reset=False):
 
   def get_mock_url_invoker(self, hard_reset=False, seed=None):
     if self.mock_url_invoker is None or hard_reset:
+      # Deferred import: tests package may not be present in production.
+      try:
+        from tests.files.mocks import MockUrlInvoker
+      except ImportError as e:
+        raise ImportError(
+            "MockUrlInvoker is only available when the 'tests' package is "
+            "present. Do not call get_mock_url_invoker() in production."
+        ) from e
+
       data_path = "tests/files/test_data/state.json"
       if seed is not None:
         data_path = f"tests/files/test_data/state_{seed}.json"
@@ -166,4 +174,4 @@ def get_files_estimator(self, progress_update_callback=lambda x: None, hard_rese
       )
       self.files_estimator.set_id_to_display_name_map(self.id_to_display_name)
     
-    return self.files_estimator 
+    return self.files_estimator
diff --git a/util/connectors.py b/util/connectors.py
index 00a73fdd..6e0ffda6 100644
--- a/util/connectors.py
+++ b/util/connectors.py
@@ -70,9 +70,17 @@ def invoke(
                 
                 final_responses += get_success_responses(responses)
                 failed_responses = get_failed_responses_that_can_be_retried(responses)
-                failed_response_ids = [response["id"] for response in failed_responses]
-
-                curr_batch = [request for request in curr_batch if str(request["id"]) in failed_response_ids]
+                # failed_responses are response dicts keyed by their "id" field.
+                # Build a set of string ids so we can match against request ids
+                # regardless of whether they were stored as int or str.
+                failed_response_ids = {str(response["id"]) for response in failed_responses}
+
+                # Rebuild curr_batch from the *original* batch list so that we
+                # keep the full request payload (not just the response slice).
+                curr_batch = [
+                    request for request in batch
+                    if str(request["id"]) in failed_response_ids
+                ]
 
                 if len(failed_responses) > 0:
                     wait_time = self.initial_delay * pow(self.batch_backoff, retry_count) + random.uniform(0, self.jitter)
@@ -82,7 +90,8 @@ def invoke(
                     break
 
             if len(failed_responses) > 0:
-                logger(f"Consistent failures observed for the following: {",".join(response.get("body") for response in failed_responses)}")
+                failed_ids = ", ".join(str(response.get("id", "?")) for response in failed_responses)
+                logger(f"Consistent failures observed for request ids: {failed_ids}")
 
         except Exception as e:
             logger(f"Error in {context}: {e}")