diff --git a/test/test_azure_helpers.py b/test/test_azure_helpers.py new file mode 100644 index 0000000000..6fe6451877 --- /dev/null +++ b/test/test_azure_helpers.py @@ -0,0 +1,155 @@ +# 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) + + +if __name__ == "__main__": + unittest.main()