fix: remove tag logic from association

Author: danabens

src/sagemaker/lineage/association.py

@@ -64,28 +64,6 @@ def delete(self):
         """Delete this Association from SageMaker."""
         self._invoke_api(self._boto_delete_method, self._boto_delete_members)
 
-    def set_tag(self, tag=None):
-        """Add a tag to the object.
-
-        Args:
-            tag (obj): Key value pair to set tag.
-
-        Returns:
-            list({str:str}): a list of key value pairs
-        """
-        return self._set_tags(resource_arn=self.source_arn, tags=[tag])
-
-    def set_tags(self, tags=None):
-        """Add tags to the object.
-
-        Args:
-            tags ([{key:value}]): list of key value pairs.
-
-        Returns:
-            list({str:str}): a list of key value pairs
-        """
-        return self._set_tags(resource_arn=self.source_arn, tags=tags)
-
     @classmethod
     def create(
         cls,

tests/integ/sagemaker/lineage/test_association.py

@@ -57,40 +57,4 @@ def test_list(association_objs, sagemaker_session):
     assert association_keys_listed
 
     for listed_asscn in listed:
-        assert listed_asscn.created_by is None
-
-
-@pytest.mark.timeout(30)
-def test_set_tag(association_obj, sagemaker_session):
-    tag = {"Key": "foo", "Value": "bar"}
-    association_obj.set_tag(tag)
-
-    while True:
-        actual_tags = sagemaker_session.sagemaker_client.list_tags(
-            ResourceArn=association_obj.source_arn
-        )["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(association_obj, sagemaker_session):
-    tags = [{"Key": "foo1", "Value": "bar1"}]
-    association_obj.set_tags(tags)
-
-    while True:
-        actual_tags = sagemaker_session.sagemaker_client.list_tags(
-            ResourceArn=association_obj.source_arn
-        )["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[-1]] == tags
+        assert listed_asscn.created_by is None

tests/unit/sagemaker/lineage/test_association.py

@@ -37,46 +37,6 @@ def test_create(sagemaker_session):
     assert "bar" == obj.destination_arn
     assert "test-type" == obj.association_type
 
-
-def test_set_tag(sagemaker_session):
-    sagemaker_session.sagemaker_client.add_association.return_value = {
-        "AssociationArn": "association-arn",
-    }
-    obj = association.Association.create(
-        source_arn="foo",
-        destination_arn="bar",
-        association_type="test-type",
-        sagemaker_session=sagemaker_session,
-    )
-    tag = {"Key": "foo", "Value": "bar"}
-    sagemaker_session.sagemaker_client.add_tags.return_value = {"Tags": tag}
-    obj.set_tag(tag)
-
-    sagemaker_session.sagemaker_client.add_tags.assert_called_with = {
-        "ResourceArn": "association-arn",
-        "Tags": [{"Key": "foo", "Value": "bar"}],
-    }
-
-
-def test_set_tags(sagemaker_session):
-    sagemaker_session.sagemaker_client.add_association.return_value = {
-        "AssociationArn": "association-arn",
-    }
-    obj = association.Association.create(
-        source_arn="foo",
-        destination_arn="bar",
-        association_type="test-type",
-        sagemaker_session=sagemaker_session,
-    )
-    tags = [{"Key": "foo1", "Value": "bar1"}, {"Key": "foo2", "Value": "bar2"}]
-    sagemaker_session.sagemaker_client.add_tags.return_value = {"Tags": tags}
-    obj.set_tags(tags=tags)
-    sagemaker_session.sagemaker_client.add_tags.assert_called_with = {
-        "ResourceArn": "association-arn",
-        "Tags": [{"Key": "foo1", "Value": "bar1"}, {"Key": "foo2", "Value": "bar2"}],
-    }
-
-
 def test_list(sagemaker_session):
     creation_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=3)