Properly propagate constructor S3 endpoint URL in S3ArtifactRepository (

mlflow#9593)

Signed-off-by: Jerry Liang <[email protected]>

mlflow/store/artifact/s3_artifact_repo.py

@@ -122,6 +122,7 @@ def _get_s3_client(self, addressing_style="path", s3_endpoint_url=None):
             access_key_id=self._access_key_id,
             secret_access_key=self._secret_access_key,
             session_token=self._session_token,
+            s3_endpoint_url=s3_endpoint_url,
         )
 
     def parse_s3_compliant_uri(self, uri):

tests/store/artifact/test_r2_artifact_repo.py

@@ -1,4 +1,6 @@
 import posixpath
+from unittest import mock
+from unittest.mock import ANY
 
 import pytest
 
@@ -40,3 +42,19 @@ def test_convert_r2_uri_to_s3_endpoint_url(r2_artifact_root):
 
     s3_endpoint_url = repo.convert_r2_uri_to_s3_endpoint_url(r2_artifact_root)
     assert s3_endpoint_url == "https://account.r2.cloudflarestorage.com"
+
+
+def test_s3_endpoint_url_is_used_to_get_s3_client(r2_artifact_root):
+    with mock.patch("boto3.client") as mock_get_s3_client:
+        artifact_uri = posixpath.join(r2_artifact_root, "some/path")
+        repo = get_artifact_repository(artifact_uri)
+        repo._get_s3_client()
+        mock_get_s3_client.assert_called_with(
+            "s3",
+            config=ANY,
+            endpoint_url="https://account.r2.cloudflarestorage.com",
+            verify=None,
+            aws_access_key_id=None,
+            aws_secret_access_key=None,
+            aws_session_token=None,
+        )