Add 30s cap time for tag tests

Author: stisac

tests/integ/sagemaker/lineage/conftest.py

@@ -155,7 +155,8 @@ def artifact_obj_with_association(sagemaker_session, artifact_obj):
 @pytest.fixture
 def trial_component_obj(sagemaker_session):
     trial_component_obj = trial_component.TrialComponent.create(
-        trial_component_name=name(), sagemaker_boto_client=sagemaker_session.sagemaker_client
+        trial_component_name=name(),
+        sagemaker_boto_client=sagemaker_session.sagemaker_client,
     )
     yield trial_component_obj
     time.sleep(0.5)
@@ -190,7 +191,9 @@ def experiment_obj(sagemaker_session):
 
 
 @pytest.fixture
-def trial_associated_artifact(artifact_obj, trial_obj, trial_component_obj, sagemaker_session):
+def trial_associated_artifact(
+    artifact_obj, trial_obj, trial_component_obj, sagemaker_session
+):
     assntn = association.Association.create(
         source_arn=artifact_obj.artifact_arn,
         destination_arn=trial_component_obj.trial_component_arn,
@@ -247,7 +250,9 @@ def artifact_obj1(sagemaker_session):
 
 
 @pytest.fixture
-def dataset_artifact_associated_models(sagemaker_session, trial_component_obj, model_artifact_obj1):
+def dataset_artifact_associated_models(
+    sagemaker_session, trial_component_obj, model_artifact_obj1
+):
     dataset_artifact_obj = artifact.DatasetArtifact.create(
         artifact_name="dataset-artifact-name",
         artifact_type="Context",
@@ -389,7 +394,9 @@ def context_obj_with_association(sagemaker_session, action_obj):
 
 
 @pytest.fixture
-def endpoint_context_associate_with_model(sagemaker_session, endpoint_action_obj, model_obj):
+def endpoint_context_associate_with_model(
+    sagemaker_session, endpoint_action_obj, model_obj
+):
     context_name = name()
     obj = context.EndpointContext.create(
         source_uri="endpontContextWithModel" + context_name,

tests/integ/sagemaker/lineage/test_action.py

@@ -16,6 +16,8 @@
 import datetime
 import logging
 
+import pytest
+
 from sagemaker.lineage import action
 
 
@@ -80,6 +82,7 @@ def test_list(action_objs, sagemaker_session):
     assert action_names
 
 
+@pytest.mark.timeout(30)
 def test_tag(action_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     action_obj.set_tag(tag)
@@ -96,6 +99,7 @@ def test_tag(action_obj, sagemaker_session):
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(action_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     action_obj.set_tags(tags)

tests/integ/sagemaker/lineage/test_artifact.py

@@ -17,6 +17,8 @@
 import logging
 import time
 
+import pytest
+
 from sagemaker.lineage import artifact
 
 
@@ -103,14 +105,17 @@ def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_sessi
     # allow trial components to index, 30 seconds max
     for i in range(3):
         time.sleep(10)
-        trials = trial_associated_artifact.downstream_trials(sagemaker_session=sagemaker_session)
+        trials = trial_associated_artifact.downstream_trials(
+            sagemaker_session=sagemaker_session
+        )
         if len(trials) > 0:
             break
 
     assert len(trials) == 1
     assert trial_obj.trial_name in trials
 
 
+@pytest.mark.timeout(30)
 def test_tag(artifact_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     artifact_obj.set_tag(tag)
@@ -127,6 +132,7 @@ def test_tag(artifact_obj, sagemaker_session):
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(artifact_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     artifact_obj.set_tags(tags)

tests/integ/sagemaker/lineage/test_association.py

@@ -16,6 +16,8 @@
 import datetime
 import time
 
+import pytest
+
 from sagemaker.lineage import association
 
 
@@ -28,7 +30,8 @@ def test_list(association_objs, sagemaker_session):
     slack = datetime.timedelta(minutes=1)
     now = datetime.datetime.now(datetime.timezone.utc)
     association_keys = [
-        assoc_obj.source_arn + ":" + assoc_obj.destination_arn for assoc_obj in association_objs
+        assoc_obj.source_arn + ":" + assoc_obj.destination_arn
+        for assoc_obj in association_objs
     ]
 
     for sort_order in ["Ascending", "Descending"]:
@@ -55,6 +58,7 @@ def test_list(association_objs, sagemaker_session):
     assert association_keys_listed
 
 
+@pytest.mark.timeout(30)
 def test_set_tag(association_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     association_obj.set_tag(tag)
@@ -72,6 +76,7 @@ def test_set_tag(association_obj, sagemaker_session):
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(association_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     association_obj.set_tags(tags)

tests/integ/sagemaker/lineage/test_context.py

@@ -16,6 +16,8 @@
 import datetime
 import logging
 
+import pytest
+
 from sagemaker.lineage import context
 
 
@@ -78,6 +80,7 @@ def test_list(context_objs, sagemaker_session):
     assert context_names
 
 
+@pytest.mark.timeout(30)
 def test_tag(context_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     context_obj.set_tag(tag)
@@ -94,6 +97,7 @@ def test_tag(context_obj, sagemaker_session):
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(context_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     context_obj.set_tags(tags)

tox.ini

@@ -55,6 +55,7 @@ markers =
     canary_quick
     cron
     local_mode
+    timeout: mark a test as a timeout.
 
 [testenv]
 passenv =