From 32f838e9adc786f086e167d5fddd4c5344471e40 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Tue, 7 Apr 2026 19:04:43 -0400 Subject: [PATCH 1/2] PYTHON-5760 Increase _azure_helpers.py coverage --- test/test_azure_helpers.py | 158 +++++++++++++++++++++++++++++++++++++ 1 file changed, 158 insertions(+) create mode 100644 test/test_azure_helpers.py diff --git a/test/test_azure_helpers.py b/test/test_azure_helpers.py new file mode 100644 index 0000000000..5ff2d1df05 --- /dev/null +++ b/test/test_azure_helpers.py @@ -0,0 +1,158 @@ +# 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 unittest.mock import patch + +sys.path[0:0] = [""] + +from pymongo._azure_helpers import _get_azure_response + + +class _MockResponse: + """Minimal context-manager response for urlopen.""" + + def __init__(self, body: str, status: int = 200): + self.status = status + self._body = body.encode("utf8") + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(body, status=400)): + 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 patch("urllib.request.urlopen", return_value=_MockResponse("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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(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 patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + self._call(timeout=42) + + _, kwargs = mock_open.call_args + self.assertEqual(kwargs["timeout"], 42) + + +if __name__ == "__main__": + unittest.main() From bdadd34fda3fd69b65fddd1a2db24cf04f31cddd Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Wed, 8 Apr 2026 17:55:02 -0400 Subject: [PATCH 2/2] Use MagicMock --- test/test_azure_helpers.py | 49 ++++++++++++++++++-------------------- 1 file changed, 23 insertions(+), 26 deletions(-) diff --git a/test/test_azure_helpers.py b/test/test_azure_helpers.py index 5ff2d1df05..6fe6451877 100644 --- a/test/test_azure_helpers.py +++ b/test/test_azure_helpers.py @@ -24,28 +24,25 @@ import json import sys import unittest -from unittest.mock import patch +from contextlib import contextmanager +from unittest.mock import MagicMock, patch sys.path[0:0] = [""] from pymongo._azure_helpers import _get_azure_response -class _MockResponse: - """Minimal context-manager response for urlopen.""" +@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") - def __init__(self, body: str, status: int = 200): - self.status = status - self._body = body.encode("utf8") - - def read(self): - return self._body - - def __enter__(self): - return self - - def __exit__(self, *args): - return False + with patch("urllib.request.urlopen", return_value=mock_response) as mock_open: + yield mock_open class TestGetAzureResponse(unittest.TestCase): @@ -54,7 +51,7 @@ def _call(self, resource="https://example.com/", client_id=None, timeout=5): def test_success_without_client_id(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + with _mock_urlopen(200, body) as mock_open: result = self._call() self.assertEqual(result["access_token"], "tok") @@ -66,7 +63,7 @@ def test_success_without_client_id(self): def test_success_with_client_id(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + with _mock_urlopen(200, body) as mock_open: result = self._call(client_id="my-client-id") self.assertEqual(result["access_token"], "tok") @@ -75,7 +72,7 @@ def test_success_with_client_id(self): def test_url_contains_resource_and_api_version(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + 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 @@ -84,7 +81,7 @@ def test_url_contains_resource_and_api_version(self): def test_request_headers(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + with _mock_urlopen(200, body) as mock_open: self._call() request = mock_open.call_args[0][0] @@ -100,14 +97,14 @@ def test_urlopen_exception_raises_value_error(self): def test_non_200_status_raises_value_error(self): body = json.dumps({"error": "something went wrong"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body, status=400)): + 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 patch("urllib.request.urlopen", return_value=_MockResponse("not-json")): + with _mock_urlopen(200, "not-json"): with self.assertRaises(ValueError) as ctx: self._call() @@ -115,7 +112,7 @@ def test_non_json_body_raises_value_error(self): def test_missing_access_token_raises_value_error(self): body = json.dumps({"expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)): + with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() @@ -123,7 +120,7 @@ def test_missing_access_token_raises_value_error(self): def test_missing_expires_in_raises_value_error(self): body = json.dumps({"access_token": "tok"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)): + with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() @@ -131,7 +128,7 @@ def test_missing_expires_in_raises_value_error(self): def test_empty_access_token_raises_value_error(self): body = json.dumps({"access_token": "", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)): + with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() @@ -139,7 +136,7 @@ def test_empty_access_token_raises_value_error(self): def test_empty_expires_in_raises_value_error(self): body = json.dumps({"access_token": "tok", "expires_in": ""}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)): + with _mock_urlopen(200, body): with self.assertRaises(ValueError) as ctx: self._call() @@ -147,7 +144,7 @@ def test_empty_expires_in_raises_value_error(self): def test_timeout_passed_to_urlopen(self): body = json.dumps({"access_token": "tok", "expires_in": "3600"}) - with patch("urllib.request.urlopen", return_value=_MockResponse(body)) as mock_open: + with _mock_urlopen(200, body) as mock_open: self._call(timeout=42) _, kwargs = mock_open.call_args