From 100d18a0d18e10d0feea2498760b12e930fe1faa Mon Sep 17 00:00:00 2001
From: Paul Van Eck <paulvaneck@microsoft.com>
Date: Fri, 22 Sep 2023 15:41:53 -0700
Subject: [PATCH] [Identity] Retry IMDS requests on 410 responses

There are cases where the IMDS endpoint might return a 410 error response,
and the guidance is to retry on such an error.

Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>
---
 sdk/identity/azure-identity/CHANGELOG.md                      | 2 ++
 .../azure-identity/azure/identity/_credentials/imds.py        | 2 +-
 sdk/identity/azure-identity/tests/test_imds_credential.py     | 4 ++--
 .../azure-identity/tests/test_imds_credential_async.py        | 4 ++--
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md
index 2bd8f91801eb..1b7ef9d87e2a 100644
--- a/sdk/identity/azure-identity/CHANGELOG.md
+++ b/sdk/identity/azure-identity/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- `ManagedIdentityCredential` will now correctly retry when the instance metadata endpoint returns a 410 response.  ([#32200](https://github.com/Azure/azure-sdk-for-python/pull/32200))
+
 ### Other Changes
 
 ## 1.15.0b1 (2023-09-12)
diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/imds.py b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py
index 23342df155bf..748a7ce12168 100644
--- a/sdk/identity/azure-identity/azure/identity/_credentials/imds.py
+++ b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py
@@ -23,7 +23,7 @@
     "connection_timeout": 2,
     "retry_backoff_factor": 2,
     "retry_backoff_max": 60,
-    "retry_on_status_codes": [404, 429] + list(range(500, 600)),
+    "retry_on_status_codes": [404, 410, 429] + list(range(500, 600)),
     "retry_status": 5,
     "retry_total": 5,
 }
diff --git a/sdk/identity/azure-identity/tests/test_imds_credential.py b/sdk/identity/azure-identity/tests/test_imds_credential.py
index d97fcff18af7..295ff297fedd 100644
--- a/sdk/identity/azure-identity/tests/test_imds_credential.py
+++ b/sdk/identity/azure-identity/tests/test_imds_credential.py
@@ -85,7 +85,7 @@ def test_imds_request_failure_docker_desktop():
 def test_retries():
     mock_response = mock.Mock(
         text=lambda encoding=None: b"{}",
-        headers={"content-type": "application/json", "Retry-After": "0"},
+        headers={"content-type": "application/json"},
         content_type="application/json",
     )
     mock_send = mock.Mock(return_value=mock_response)
@@ -93,7 +93,7 @@ def test_retries():
     total_retries = PIPELINE_SETTINGS["retry_total"]
 
     assert within_credential_chain.get() == False
-    for status_code in (404, 429, 500):
+    for status_code in (404, 410, 429, 500):
         mock_send.reset_mock()
         mock_response.status_code = status_code
         try:
diff --git a/sdk/identity/azure-identity/tests/test_imds_credential_async.py b/sdk/identity/azure-identity/tests/test_imds_credential_async.py
index ae4c737df027..948f5f20f34f 100644
--- a/sdk/identity/azure-identity/tests/test_imds_credential_async.py
+++ b/sdk/identity/azure-identity/tests/test_imds_credential_async.py
@@ -160,14 +160,14 @@ async def test_cache():
 async def test_retries():
     mock_response = mock.Mock(
         text=lambda encoding=None: b"{}",
-        headers={"content-type": "application/json", "Retry-After": "0"},
+        headers={"content-type": "application/json"},
         content_type="application/json",
     )
     mock_send = mock.Mock(return_value=mock_response)
 
     total_retries = PIPELINE_SETTINGS["retry_total"]
 
-    for status_code in (404, 429, 500):
+    for status_code in (404, 410, 429, 500):
         mock_send.reset_mock()
         mock_response.status_code = status_code
         try: