sleep time before clean_all added (avoid race condition)

Author: Yi-Ting Lee

tests/integ/sagemaker/lineage/helpers.py

@@ -142,6 +142,10 @@ def create_association(self, source_arn, dest_arn, association_type="AssociatedW
             return False
 
     def clean_all(self):
+        # clean all lineage data created by LineageResourceHelper
+
+        time.sleep(1)  # avoid GSI race condition between create & delete
+
         for source, dest in self.associations:
             try:
                 self.client.delete_association(SourceArn=source, DestinationArn=dest)

tests/integ/sagemaker/lineage/test_lineage_visualize.py

@@ -29,10 +29,11 @@ def test_LineageResourceHelper(sagemaker_session):
         art1 = lineage_resource_helper.create_artifact(artifact_name=name())
         art2 = lineage_resource_helper.create_artifact(artifact_name=name())
         lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
-        lineage_resource_helper.clean_all()
     except Exception as e:
         print(e)
         assert False
+    finally:
+        lineage_resource_helper.clean_all()
 
 
 @pytest.mark.skip("visualizer load test")
@@ -46,7 +47,7 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(200):
+        for i in range(150):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(
                 source_arn=wide_graph_root_arn, dest_arn=artifact_arn
@@ -56,6 +57,10 @@ def test_wide_graph_visualize(sagemaker_session):
         lq_result = lq.query(start_arns=[wide_graph_root_arn])
         lq_result.visualize(path="wideGraph.html")
 
+        print("vertex len = ")
+        print(len(lq_result.vertices))
+        assert False
+
     except Exception as e:
         print(e)
         assert False