2 files changed, 129 insertions, 1 deletions
diff --git a/tests/unittests/sources/azure/test_errors.py b/tests/unittests/sources/azure/test_errors.py
index e9d3e039..cf3e0e71 100644
--- a/tests/unittests/sources/azure/test_errors.py
+++ b/tests/unittests/sources/azure/test_errors.py
@@ -5,9 +5,11 @@ import datetime
 from unittest import mock
 
 import pytest
+import requests
 
 from cloudinit import version
 from cloudinit.sources.azure import errors
+from cloudinit.url_helper import UrlError
 
 
 @pytest.fixture()
@@ -134,6 +136,69 @@ def test_dhcp_interface_not_found():
     assert error.supporting_data["duration"] == 5.6
 
 
+@pytest.mark.parametrize(
+    "exception,reason",
+    [
+        (
+            UrlError(
+                requests.ConnectionError(),
+            ),
+            "connection error querying IMDS",
+        ),
+        (
+            UrlError(
+                requests.ConnectTimeout(),
+            ),
+            "connection timeout querying IMDS",
+        ),
+        (
+            UrlError(
+                requests.ReadTimeout(),
+            ),
+            "read timeout querying IMDS",
+        ),
+        (
+            UrlError(
+                Exception(),
+                code=500,
+            ),
+            "http error querying IMDS",
+        ),
+        (
+            UrlError(
+                requests.HTTPError(),
+                code=None,
+            ),
+            "unexpected error querying IMDS",
+        ),
+    ],
+)
+def test_imds_url_error(exception, reason):
+    duration = 123.4
+    fake_url = "fake://url"
+
+    exception.url = fake_url
+    error = errors.ReportableErrorImdsUrlError(
+        exception=exception, duration=duration
+    )
+
+    assert error.reason == reason
+    assert error.supporting_data["duration"] == duration
+    assert error.supporting_data["exception"] == repr(exception)
+    assert error.supporting_data["url"] == fake_url
+
+
+def test_imds_metadata_parsing_exception():
+    exception = ValueError("foobar")
+
+    error = errors.ReportableErrorImdsMetadataParsingException(
+        exception=exception
+    )
+
+    assert error.reason == "error parsing IMDS metadata"
+    assert error.supporting_data["exception"] == repr(exception)
+
+
 def test_unhandled_exception():
     source_error = None
     try:
diff --git a/tests/unittests/sources/test_azure.py b/tests/unittests/sources/test_azure.py
index 59ac6459..ed605ee7 100644
--- a/tests/unittests/sources/test_azure.py
+++ b/tests/unittests/sources/test_azure.py
@@ -17,7 +17,7 @@ from cloudinit.reporting.handlers import HyperVKvpReportingHandler
 from cloudinit.sources import UNSET
 from cloudinit.sources import DataSourceAzure as dsaz
 from cloudinit.sources import InvalidMetaDataException
-from cloudinit.sources.azure import errors, identity
+from cloudinit.sources.azure import errors, identity, imds
 from cloudinit.sources.helpers import netlink
 from cloudinit.util import (
     MountFailedError,
@@ -136,6 +136,16 @@ def mock_ephemeral_dhcp_v4():
 
 
 @pytest.fixture
+def mock_imds_fetch_metadata_with_api_fallback():
+    with mock.patch.object(
+        imds,
+        "fetch_metadata_with_api_fallback",
+        autospec=True,
+    ) as m:
+        yield m
+
+
+@pytest.fixture
 def mock_kvp_report_failure_to_host():
     with mock.patch(
         MOCKPATH + "kvp.report_failure_to_host",
@@ -4244,6 +4254,59 @@ class TestProvisioning:
         assert len(self.mock_kvp_report_success_to_host.mock_calls) == 1
 
 
+class TestGetMetadataFromImds:
+    @pytest.mark.parametrize("report_failure", [False, True])
+    @pytest.mark.parametrize(
+        "exception,reported_error_type",
+        [
+            (
+                url_helper.UrlError(requests.ConnectionError()),
+                errors.ReportableErrorImdsUrlError,
+            ),
+            (
+                ValueError("bad data"),
+                errors.ReportableErrorImdsMetadataParsingException,
+            ),
+        ],
+    )
+    def test_errors(
+        self,
+        azure_ds,
+        exception,
+        mock_azure_report_failure_to_fabric,
+        mock_imds_fetch_metadata_with_api_fallback,
+        mock_kvp_report_failure_to_host,
+        monkeypatch,
+        report_failure,
+        reported_error_type,
+    ):
+        monkeypatch.setattr(
+            azure_ds, "_is_ephemeral_networking_up", lambda: True
+        )
+        mock_imds_fetch_metadata_with_api_fallback.side_effect = exception
+
+        assert (
+            azure_ds.get_metadata_from_imds(report_failure=report_failure)
+            == {}
+        )
+        assert mock_imds_fetch_metadata_with_api_fallback.mock_calls == [
+            mock.call(retry_deadline=mock.ANY)
+        ]
+
+        reported_error = mock_kvp_report_failure_to_host.call_args[0][0]
+        assert isinstance(reported_error, reported_error_type)
+        assert reported_error.supporting_data["exception"] == repr(exception)
+        assert mock_kvp_report_failure_to_host.mock_calls == [
+            mock.call(reported_error)
+        ]
+        if report_failure:
+            assert mock_azure_report_failure_to_fabric.mock_calls == [
+                mock.call(endpoint=mock.ANY, error=reported_error)
+            ]
+        else:
+            assert mock_azure_report_failure_to_fabric.mock_calls == []
+
+
 class TestReportFailure:
     @pytest.mark.parametrize("kvp_enabled", [False, True])
     def report_host_only_kvp_enabled(