From e6f9b7aa98fc13f8f72ec233e356a1ef6b40290d Mon Sep 17 00:00:00 2001
From: Iain-S <25081046+Iain-S@users.noreply.github.com>
Date: Thu, 7 Sep 2023 09:19:37 +0100
Subject: [PATCH] Ignore resources when disabling subscriptions

---
 .../controller/subscription.py                | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/controller_function/controller/subscription.py b/controller_function/controller/subscription.py
index 6666d1a..f803ef6 100644
--- a/controller_function/controller/subscription.py
+++ b/controller_function/controller/subscription.py
@@ -1,16 +1,45 @@
 """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(
+    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.
 
@@ -53,3 +82,16 @@ 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