[TP-1356] Wait for inputs to be deleted

tests/client/test_annotations_searches.py

@@ -5,6 +5,7 @@
     wait_for_inputs_upload,
     metadata,
     both_channels,
+    cleanup_inputs,
     TRAVEL_IMAGE_URL,
 )
 
@@ -204,7 +205,4 @@ def __enter__(self) -> resources_pb2.Input:
         return self._input
 
     def __exit__(self, type_, value, traceback) -> None:
-        delete_response = self._stub.DeleteInput(
-            service_pb2.DeleteInputRequest(input_id=self._input.id), metadata=metadata()
-        )
-        raise_on_failure(delete_response)
+        cleanup_inputs(self._stub, [self._input.id], metadata=metadata())

tests/client/test_dataset_version_export.py

@@ -11,6 +11,7 @@
     both_channels,
     metadata,
     raise_on_failure,
+    wait_for_inputs_delete,
     wait_for_inputs_upload,
     wait_for_dataset_version_ready,
     wait_for_dataset_version_export_success,
@@ -162,6 +163,7 @@ def check_json(batch_str):
             raise_on_failure(delete_dataset_versions_response)
         if input_ids:
             raise_on_failure(delete_inputs_response)
+            wait_for_inputs_delete(stub, input_ids, metadata=metadata())
         raise_on_failure(delete_datasets_response)
 
 

tests/client/test_image_crud.py

@@ -7,6 +7,7 @@
     RED_TRUCK_IMAGE_FILE_PATH,
     TRUCK_IMAGE_URL,
     both_channels,
+    cleanup_inputs,
     metadata,
     raise_on_failure,
     wait_for_inputs_upload,
@@ -71,9 +72,7 @@ def test_post_list_patch_get_delete_image(channel):
         raise_on_failure(get_response)
         assert get_response.input.data.concepts[0].name == "some-new-concept"
     finally:
-        delete_request = service_pb2.DeleteInputRequest(input_id=input_id)
-        delete_response = stub.DeleteInput(delete_request, metadata=metadata())
-        raise_on_failure(delete_response)
+        cleanup_inputs(stub, [input_id], metadata=metadata())
 
 
 @both_channels
@@ -103,10 +102,7 @@ def test_post_delete_batch_images(channel):
 
     wait_for_inputs_upload(stub, metadata(), [input_id1, input_id2])
 
-    delete_response = stub.DeleteInputs(
-        service_pb2.DeleteInputsRequest(ids=[input_id1, input_id2]), metadata=metadata()
-    )
-    raise_on_failure(delete_response)
+    cleanup_inputs(stub, [input_id1, input_id2], metadata=metadata())
 
 
 @both_channels
@@ -198,10 +194,7 @@ def test_post_patch_get_image_with_id_concepts_geo_and_metadata(channel):
         inp.data.metadata["key1"] == 123
     )  # Since we use the merge action, the old values should remain
 
-    delete_response = stub.DeleteInputs(
-        service_pb2.DeleteInputsRequest(ids=[input_id]), metadata=metadata()
-    )
-    raise_on_failure(delete_response)
+    cleanup_inputs(stub, [input_id], metadata=metadata())
 
 
 @both_channels
@@ -225,8 +218,4 @@ def test_image_with_bytes(channel):
     input_id = post_response.inputs[0].id
 
     wait_for_inputs_upload(stub, metadata(), [input_id])
-
-    delete_response = stub.DeleteInputs(
-        service_pb2.DeleteInputsRequest(ids=[input_id]), metadata=metadata()
-    )
-    raise_on_failure(delete_response)
+    cleanup_inputs(stub, [input_id], metadata=metadata())

tests/client/test_inputs_searches.py

@@ -17,6 +17,7 @@
 from clarifai_grpc.grpc.api.service_pb2 import PostInputsSearchesRequest
 from tests.common import (
     both_channels,
+    cleanup_inputs,
     metadata,
     raise_on_failure,
     DOG_IMAGE_URL,
@@ -433,7 +434,4 @@ def __enter__(self) -> resources_pb2.Input:
         return self._input
 
     def __exit__(self, type_, value, traceback) -> None:
-        delete_response = self._stub.DeleteInput(
-            service_pb2.DeleteInputRequest(input_id=self._input.id), metadata=metadata()
-        )
-        raise_on_failure(delete_response)
+        cleanup_inputs(self._stub, [self._input.id], metadata=metadata())

tests/client/test_model_crud.py

@@ -6,6 +6,7 @@
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2, service_pb2_grpc
 from tests.common import (
     both_channels,
+    cleanup_inputs,
     metadata,
     raise_on_failure,
     wait_for_model_evaluated,
@@ -186,11 +187,7 @@ def test_post_patch_get_train_evaluate_predict_delete_model(channel):
         )
         raise_on_failure(delete_response)
 
-        delete_inputs_response = stub.DeleteInputs(
-            service_pb2.DeleteInputsRequest(ids=[input_id_1, input_id_2]),
-            metadata=metadata(),
-        )
-        raise_on_failure(delete_inputs_response)
+        cleanup_inputs(stub, [input_id_1, input_id_2], metadata=metadata())
 
 
 @both_channels
@@ -359,6 +356,4 @@ def test_model_creation_training_and_evaluation(channel):
         stub.DeleteModel(service_pb2.DeleteModelRequest(model_id=model_id), metadata=metadata())
     )
 
-    raise_on_failure(
-        stub.DeleteInputs(service_pb2.DeleteInputsRequest(ids=input_ids), metadata=metadata())
-    )
+    cleanup_inputs(stub, input_ids, metadata=metadata())

tests/client/test_secure_data_hosting.py

@@ -9,6 +9,7 @@
     TRAVEL_IMAGE_URL,
     BEER_VIDEO_URL,
     both_channels,
+    cleanup_inputs,
     raise_on_failure,
     wait_for_inputs_upload,
 )
@@ -236,8 +237,4 @@ def test_adding_inputs(channel):
                 verify_url_with_all_auths(input_url_from_get, verify_func=fn)
                 verify_url_with_bad_auth(input_url_from_get)  # these should fail
     finally:
-        # delete inputs
-        for cfid in bytes_hash_by_id.keys():
-            delete_request = service_pb2.DeleteInputRequest(input_id=cfid)
-            delete_response = stub.DeleteInput(delete_request, metadata=API_CLIENT_AUTH)
-            raise_on_failure(delete_response)
+        cleanup_inputs(stub, bytes_hash_by_id.keys(), metadata=API_CLIENT_AUTH)

tests/common.py

@@ -85,6 +85,26 @@ def wait_for_inputs_upload(stub, metadata, input_ids):
                 )
     # At this point, all inputs have been downloaded successfully.
 
+def cleanup_inputs(stub, input_ids, metadata):
+    delete_request = service_pb2.DeleteInputsRequest(ids=input_ids)
+    delete_response = stub.DeleteInputs(delete_request, metadata=metadata)
+    raise_on_failure(delete_response)
+    wait_for_inputs_delete(stub, metadata, input_ids)
+
+def wait_for_inputs_delete(stub, input_ids, metadata):
+    remaining_input_ids = list(input_ids)
+    while remaining_input_ids:
+        for input_id in remaining_input_ids:
+            get_input_response = stub.GetInput(
+                service_pb2.GetInputRequest(input_id=input_id), metadata=metadata
+            )
+            if get_input_response.status.code == status_code_pb2.CONN_DOES_NOT_EXIST:
+                remaining_input_ids.remove(input_id)
+            else:
+                print(f"Waiting for input '{input_id}' to be deleted")
+                time.sleep(1)
+                break
+
 
 def wait_for_model_trained(stub, metadata, model_id, model_version_id, user_app_id=None):
     while True: