From 44a0d305da8741274e9d8f934af5a3c4fe2ddb70 Mon Sep 17 00:00:00 2001
From: Iain Stenson <25081046+Iain-S@users.noreply.github.com>
Date: Thu, 4 Jul 2024 09:17:27 +0100
Subject: [PATCH] Add IngoreResourceCheck to new Azure API call (#32)

---
 .../controller/subscription.py                | 50 +++----------------
 1 file changed, 7 insertions(+), 43 deletions(-)

diff --git a/controller_function/controller/subscription.py b/controller_function/controller/subscription.py
index d54d88d..7cfc5be 100644
--- a/controller_function/controller/subscription.py
+++ b/controller_function/controller/subscription.py
@@ -1,45 +1,16 @@
 """Manage subscription life cycle."""
 import logging
-import types
 from uuid import UUID
 
 from azure.core import exceptions as azure_exceptions
 from azure.identity import DefaultAzureCredential
 from azure.mgmt.subscription import SubscriptionClient
 
-# pylint: disable=too-many-arguments
-
 CREDENTIALS = DefaultAzureCredential()
 
 SUBSCRIPTION_CLIENT = SubscriptionClient(credential=CREDENTIALS)
 
 
-def new_post(  # type: ignore
-    self,
-    url,
-    params=None,
-    headers=None,
-    content=None,
-    form_content=None,
-    stream_content=None,
-):
-    """Add IgnoreResourceCheck=True to the existing query params.
-
-    See also https://github.com/Azure/azure-sdk-for-python/issues/10814
-    """
-    if url.endswith("cancel"):
-        params = {**(params or {}), **{"IgnoreResourceCheck": True}}
-
-    return self.original_post_method(
-        url,
-        params=params,
-        headers=headers,
-        content=content,
-        form_content=form_content,
-        stream_content=stream_content,
-    )
-
-
 def enable_subscription(subscription_id: UUID) -> None:
     """Enable a subscription.
 
@@ -72,7 +43,13 @@ def disable_subscription(subscription_id: UUID) -> None:
         subscription_id: Subscription id.
     """
     try:
-        SUBSCRIPTION_CLIENT.subscription.cancel(str(subscription_id))
+        # IgnoreResourceCheck=True is required to disable a subscription
+        # that has resources.
+        # See https://github.com/Azure/azure-sdk-for-python/issues/10814
+        SUBSCRIPTION_CLIENT.subscription.cancel(
+            str(subscription_id),
+            params={"IgnoreResourceCheck": True},
+        )
     except azure_exceptions.HttpResponseError as e:
         is_disabled = (
             ("Subscription is not in active state" in e.error.message)
@@ -87,16 +64,3 @@ def disable_subscription(subscription_id: UUID) -> None:
             )
         else:
             raise e
-
-
-# Patch the original post method with our own so that we can add
-# IgnoreResourceCheck=True to the query params. Without this,
-# subscriptions with resources will not be disabled.
-# pylint: disable=protected-access
-SUBSCRIPTION_CLIENT.subscription._client.original_post_method = (
-    SUBSCRIPTION_CLIENT.subscription._client.post
-)
-SUBSCRIPTION_CLIENT.subscription._client.post = types.MethodType(
-    new_post, SUBSCRIPTION_CLIENT.subscription._client
-)
-# pylint: enable=protected-access