# Copyright 2026-present MongoDB, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Unit tests for _azure_helpers.py. These tests mock urlopen to avoid requiring a live Azure IMDS endpoint. Integration tests that exercise the real endpoint are gated by environment variables in test_on_demand_csfle.py and test_auth_oidc.py. """ from __future__ import annotations import json import sys import unittest from contextlib import contextmanager from unittest.mock import MagicMock, patch sys.path[0:0] = [""] from pymongo._azure_helpers import _get_azure_response @contextmanager def _mock_urlopen(status: int, body: str): """Context manager that patches ``urllib.request.urlopen`` with a fake response.""" mock_response = MagicMock() mock_response.__enter__ = lambda s: s mock_response.__exit__ = MagicMock(return_value=False) mock_response.status = status mock_response.read.return_value = body.encode("utf8") with patch("urllib.request.urlopen", return_value=mock_response) as mock_open: yield mock_open class TestGetAzureResponse(unittest.TestCase): def _call(self, resource="https://example.com/", client_id=None, timeout=5): return _get_azure_response(resource, client_id=client_id, timeout=timeout) def test_success_without_client_id(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: result = self._call() self.assertEqual(result["access_token"], "tok") self.assertEqual(result["expires_in"], "3600") # Verify client_id was NOT added to the URL url = mock_open.call_args[0][0].full_url self.assertNotIn("client_id", url) def test_success_with_client_id(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: result = self._call(client_id="my-client-id") self.assertEqual(result["access_token"], "tok") url = mock_open.call_args[0][0].full_url self.assertIn("client_id=my-client-id", url) def test_url_contains_resource_and_api_version(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: self._call(resource="https://test-resource.example.com") url = mock_open.call_args[0][0].full_url self.assertIn("api-version=2018-02-01", url) self.assertIn("resource=https://test-resource.example.com", url) def test_request_headers(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: self._call() request = mock_open.call_args[0][0] self.assertEqual(request.get_header("Metadata"), "true") self.assertEqual(request.get_header("Accept"), "application/json") def test_urlopen_exception_raises_value_error(self): with patch("urllib.request.urlopen", side_effect=OSError("connection refused")): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("Failed to acquire IMDS access token", str(ctx.exception)) def test_non_200_status_raises_value_error(self): body = json.dumps({"error": "something went wrong"}) with _mock_urlopen(400, body): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("Failed to acquire IMDS access token", str(ctx.exception)) def test_non_json_body_raises_value_error(self): with _mock_urlopen(200, "not-json"): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("Azure IMDS response must be in JSON format", str(ctx.exception)) def test_missing_access_token_raises_value_error(self): body = json.dumps({"expires_in": "3600"}) with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("access_token", str(ctx.exception)) def test_missing_expires_in_raises_value_error(self): body = json.dumps({"access_token": "tok"}) with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("expires_in", str(ctx.exception)) def test_empty_access_token_raises_value_error(self): body = json.dumps({"access_token": "", "expires_in": "3600"}) with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("access_token", str(ctx.exception)) def test_empty_expires_in_raises_value_error(self): body = json.dumps({"access_token": "tok", "expires_in": ""}) with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() self.assertIn("expires_in", str(ctx.exception)) def test_timeout_passed_to_urlopen(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: self._call(timeout=42) _, kwargs = mock_open.call_args self.assertEqual(kwargs["timeout"], 42) def test_client_id_is_url_encoded(self): """Ensure special characters in client_id are percent-encoded.""" body = json.dumps({"access_token": "tok", "expires_in": "3600"}) with _mock_urlopen(200, body) as mock_open: self._call(client_id="id with spaces&special=chars") url = mock_open.call_args[0][0].full_url # '&' and '=' must be percent-encoded so they don't inject extra query params self.assertIn("client_id=id%20with%20spaces%26special%3Dchars", url) # The encoded client_id should not introduce a raw '&' # Count params: api-version, resource, client_id — exactly 3 query_string = url.split("?", 1)[1] self.assertEqual(query_string.count("&"), 2) if __name__ == "__main__": unittest.main()