Add 30s cap time for tag tests

setup.py

@@ -66,6 +66,7 @@ def read_version():
         "pytest<6.1.0",
         "pytest-cov",
         "pytest-rerunfailures",
+        "pytest-timeout",
         "pytest-xdist",
         "mock",
         "contextlib2",

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)

tests/integ/sagemaker/lineage/test_action.py

@@ -15,6 +15,9 @@
 
 import datetime
 import logging
+import time
+
+import pytest
 
 from sagemaker.lineage import action
 
@@ -80,6 +83,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)
@@ -90,12 +94,14 @@ def test_tag(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     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)
@@ -106,6 +112,7 @@ def test_tags(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_artifact.py

@@ -17,6 +17,8 @@
 import logging
 import time
 
+import pytest
+
 from sagemaker.lineage import artifact
 
 
@@ -111,6 +113,7 @@ def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_sessi
     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)
@@ -121,12 +124,14 @@ def test_tag(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     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)
@@ -137,6 +142,7 @@ def test_tags(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_association.py

@@ -16,6 +16,8 @@
 import datetime
 import time
 
+import pytest
+
 from sagemaker.lineage import association
 
 
@@ -55,6 +57,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)
@@ -65,13 +68,14 @@ def test_set_tag(association_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-        time.sleep(1)
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     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)
@@ -82,7 +86,7 @@ def test_tags(association_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-        time.sleep(1)
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_context.py

@@ -15,6 +15,9 @@
 
 import datetime
 import logging
+import time
+
+import pytest
 
 from sagemaker.lineage import context
 
@@ -78,6 +81,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)
@@ -88,12 +92,14 @@ def test_tag(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     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)
@@ -104,6 +110,7 @@ def test_tags(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tox.ini

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