Updated existing tests

Author: dhanur-sharma

inference/tests/local_test_inference_api_client.py

@@ -1,4 +1,5 @@
-# inference/tests/test_inference_api_client.py
+# inference/tests/local_test_inference_api_client.py
+# docker-compose -f local.yml run --rm django pytest inference/tests/local_test_inference_api_client.py
 
 """
 This is a test designed to be run on a local machine which has the inference pipeline running

inference/tests/test_batch.py

@@ -1,4 +1,5 @@
 # inference/tests/test_batch.py
+# docker-compose -f local.yml run --rm django pytest inference/tests/test_batch.py
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest

inference/tests/test_inference_integration.py

@@ -1,5 +1,7 @@
 # inference/tests/test_inference_integration.py
 # docker-compose -f local.yml run --rm django pytest inference/tests/test_inference_integration.py
+import time
+
 import pytest
 
 from inference.models.inference import ExternalJob, InferenceJob, ModelVersion
@@ -152,6 +154,17 @@ def test_create_external_job(loaded_model, inference_job, api_client):
     external_job.refresh_from_db()
 
     # The job should eventually complete (this might take time)
+    max_retries = 10
+    retry_count = 0
+    while retry_count < max_retries:
+        external_job.refresh_status_and_store_results()
+        external_job.refresh_from_db()
+
+        if external_job.status in [ExternalJobStatus.COMPLETED, ExternalJobStatus.FAILED, ExternalJobStatus.CANCELLED]:
+            break
+
+        retry_count += 1
+        time.sleep(1)
     assert external_job.status in [ExternalJobStatus.COMPLETED, ExternalJobStatus.FAILED, ExternalJobStatus.CANCELLED]
 
     # If job completed, check for results