Language: Python
Label: safe
CWE: CWE-918 \u2014 Server-Side Request Forgery (SSRF)
The web server receives a URL or similar request from an upstream component and retrieves the contents of this URL, but it does not sufficiently ensure that the request is being sent to the expected destination.

Code:
def test_log_artifact_adls_gen2_flush_error(databricks_artifact_repo, test_file):
    mock_successful_response = Response()
    mock_successful_response.status_code = 200
    mock_successful_response.close = lambda: None
    mock_error_response = MlflowException("MOCK ERROR")
    mock_credential_info = ArtifactCredentialInfo(
        signed_uri=MOCK_ADLS_GEN2_SIGNED_URI,
        type=ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI,
    )
    with mock.patch(
        f"{DATABRICKS_ARTIFACT_REPOSITORY}._get_credential_infos",
        return_value=[mock_credential_info],
    ) as get_credential_infos_mock, mock.patch(
        "requests.Session.request", side_effect=[mock_successful_response, mock_error_response]
    ) as request_mock:
        mock_credential_info = ArtifactCredentialInfo(
            signed_uri=MOCK_ADLS_GEN2_SIGNED_URI,
            type=ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI,
        )
        with pytest.raises(MlflowException, match=r"MOCK ERROR"):
            databricks_artifact_repo.log_artifact(test_file)
        get_credential_infos_mock.assert_called_with(GetCredentialsForWrite, MOCK_RUN_ID, ANY)
        assert request_mock.mock_calls == [
            mock.call(
                "put",
                f"{MOCK_ADLS_GEN2_SIGNED_URI}?resource=file",
                allow_redirects=True,
                headers={},
                timeout=None,
            ),
            mock.call(
                "patch",
                f"{MOCK_ADLS_GEN2_SIGNED_URI}?action=append&position=0&flush=true",
                allow_redirects=True,
                data=ANY,
                headers={},
                timeout=None,
            ),
        ]