Fix unit tests failures

requirements/extras/test_requirements.txt

@@ -1,4 +1,5 @@
 tox==3.24.5
+setuptools
 numpy>=2.0.0, <3.0
 build[virtualenv]==1.2.1
 flake8==7.1.2

tests/unit/sagemaker/feature_store/test_dataset_builder.py

@@ -31,7 +31,9 @@
 
 @pytest.fixture
 def sagemaker_session_mock():
-    return Mock()
+    mock = Mock()
+    mock.sagemaker_config = None
+    return mock
 
 
 @pytest.fixture

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -68,7 +68,7 @@ def s3_uri():
 @pytest.fixture
 def sagemaker_session_mock():
     sagemaker_session_mock = Mock()
-    sagemaker_session_mock.sagemaker_config = {}
+    sagemaker_session_mock.sagemaker_config = None
     return sagemaker_session_mock
 
 

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -33,7 +33,9 @@
 
 @pytest.fixture
 def sagemaker_session_mock():
-    return Mock()
+    mock = Mock()
+    mock.sagemaker_config = None
+    return mock
 
 
 @pytest.fixture

tests/unit/sagemaker/feature_store/test_feature_utils.py

@@ -42,7 +42,9 @@ def __reduce__(self):
 @pytest.fixture
 def sagemaker_session_mock():
     """Fixture Mock class"""
-    return Mock()
+    mock = Mock()
+    mock.sagemaker_config = None
+    return mock
 
 
 def test_convert_unsupported_types_to_supported(sagemaker_session_mock):

tests/unit/sagemaker/lineage/test_artifact.py

@@ -323,7 +323,12 @@ def test_create_delete_with_association(sagemaker_session):
     )
 
 
-def test_downstream_trials(sagemaker_session):
+@unittest.mock.patch("sagemaker.lineage.artifact.get_module")
+def test_downstream_trials(mock_get_module, sagemaker_session):
+    # Mock smexperiments module
+    mock_smexperiments = unittest.mock.MagicMock()
+    mock_get_module.return_value = mock_smexperiments
+
     sagemaker_session.sagemaker_client.list_associations.side_effect = [
         {
             "AssociationSummaries": [
@@ -379,7 +384,12 @@ def test_downstream_trials(sagemaker_session):
     assert expected_calls == sagemaker_session.sagemaker_client.list_associations.mock_calls
 
 
-def test_downstream_trials_v2(sagemaker_session):
+@unittest.mock.patch("sagemaker.lineage.artifact.get_module")
+def test_downstream_trials_v2(mock_get_module, sagemaker_session):
+    # Mock smexperiments module
+    mock_smexperiments = unittest.mock.MagicMock()
+    mock_get_module.return_value = mock_smexperiments
+
     sagemaker_session.sagemaker_client.query_lineage.return_value = {
         "Vertices": [
             {"Arn": "B" + str(i), "Type": "DataSet", "LineageType": "Artifact"} for i in range(10)
@@ -425,7 +435,12 @@ def test_downstream_trials_v2(sagemaker_session):
     assert expected_calls == sagemaker_session.sagemaker_client.query_lineage.mock_calls
 
 
-def test_upstream_trials(sagemaker_session):
+@unittest.mock.patch("sagemaker.lineage.artifact.get_module")
+def test_upstream_trials(mock_get_module, sagemaker_session):
+    # Mock smexperiments module
+    mock_smexperiments = unittest.mock.MagicMock()
+    mock_get_module.return_value = mock_smexperiments
+
     sagemaker_session.sagemaker_client.query_lineage.return_value = {
         "Vertices": [
             {"Arn": "B" + str(i), "Type": "DataSet", "LineageType": "Artifact"} for i in range(10)

tests/unit/sagemaker/workflow/test_callback_step.py

@@ -28,6 +28,7 @@
 def sagemaker_session_mock():
     session_mock = Mock()
     session_mock.boto_session.client = Mock()
+    session_mock.sagemaker_config = None
     return session_mock
 
 

tests/unit/test_mxnet.py

@@ -19,7 +19,7 @@
 import pytest
 from mock import MagicMock, Mock, ANY
 from mock import patch
-from pkg_resources import parse_version
+from packaging.version import parse as parse_version
 
 from sagemaker.fw_utils import UploadedCode
 from sagemaker.drift_check_baselines import DriftCheckBaselines

tox.ini

@@ -94,6 +94,7 @@ commands =
 
     pytest {posargs}
 deps =
+    setuptools
     .[test]
     asyncio
     nest_asyncio