[TESTS] Fix retry mechanism to raise if retry didn't help

mvafin · Oct 18, 2024 · 071ed86 · 071ed86
1 parent 939b35a
commit 071ed86
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/tests/layer_tests/py_frontend_tests/test_torchvision_preprocessor.py b/tests/layer_tests/py_frontend_tests/test_torchvision_preprocessor.py
@@ -32,15 +32,17 @@ def forward(self, data):
 def _infer_pipelines(test_input, preprocess_pipeline, input_channels=3):
     retries = 0
     max_retries = 3
+    last_e = None
     while retries < max_retries:
         try:
             return _infer_pipelines_impl(test_input, preprocess_pipeline, input_channels)
         except RuntimeError as e:
             # This is a potentially sporadic issue
             print(f"An error occurred: {e}. Retrying...")
+            last_e = e
             retries += 1
     else:
-        print("Max retries reached. Function execution failed.")
+        raise RuntimeError("Max retries reached. Function execution failed.") from last_e
 
 
 def _infer_pipelines_impl(test_input, preprocess_pipeline, input_channels=3):

diff --git a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py
@@ -73,19 +73,21 @@ def _test(self, model, ref_net, kind, ie_device, precision, ir_version, infer_ti
               **kwargs):
         retries = 0
         max_retries = 3
+        last_e = None
         while retries < max_retries:
             try:
                 return self._test_impl(model, ref_net, kind, ie_device, precision, ir_version, infer_timeout, dynamic_shapes, **kwargs)
             except RuntimeError as e:
                 # This is a potentially sporadic issue
                 print(f"An error occurred: {e}. Retrying...")
+                last_e = e
                 retries += 1
         else:
-            print("Max retries reached. Function execution failed.")
+            raise RuntimeError("Max retries reached. Function execution failed.") from last_e
 
 
     def _test_impl(self, model, ref_net, kind, ie_device, precision, ir_version, infer_timeout=60, dynamic_shapes=True,
-              **kwargs):
+                   **kwargs):
         """
         :param enabled_transforms/disabled_transforms: string with idxs of transforms that should be enabled/disabled.
                                                        Example: "transform_1,transform_2"