From 4a4969936f8f53e919b3c203e53b9d952df53f1a Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 11:49:20 -0400
Subject: [PATCH 1/8] Change the input selection process

---
 shared/python/utils.py | 81 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 70 insertions(+), 11 deletions(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index abed46e..889cb55 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -266,12 +266,13 @@ def __init__(self, rg_location: str, deployment: INFRASTRUCTURE, index: int, api
     #    PUBLIC METHODS
     # ------------------------------
 
-    def create_infrastructure(self, bypass_infrastructure_check: bool = False) -> bool:
+    def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow_update: bool = True) -> bool:
         """
         Create infrastructure by executing the appropriate creation script.
         
         Args:
             bypass_infrastructure_check (bool): Skip infrastructure existence check. Defaults to False.
+            allow_update (bool): Allow infrastructure updates when infrastructure already exists. Defaults to True.
             
         Returns:
             bool: True if infrastructure creation succeeded, False otherwise.
@@ -279,7 +280,14 @@ def create_infrastructure(self, bypass_infrastructure_check: bool = False) -> bo
 
         import sys 
 
-        if bypass_infrastructure_check or not does_infrastructure_exist(self.deployment, self.index):
+        try:
+            # For infrastructure notebooks, allow the update option
+            infrastructure_exists = does_infrastructure_exist(self.deployment, self.index, allow_update_option=allow_update)
+        except SystemExit:
+            # User cancelled the operation
+            return False
+        
+        if bypass_infrastructure_check or not infrastructure_exists:
             # Map infrastructure types to their folder names
             infra_folder_map = {
                 INFRASTRUCTURE.SIMPLE_APIM: 'simple-apim',
@@ -537,9 +545,16 @@ def _query_and_select_infrastructure(self) -> tuple[INFRASTRUCTURE | None, int |
                     
             except ValueError:
                 print_error('Invalid input. Please enter a number.')
-            except KeyboardInterrupt:
-                print_warning('\nOperation cancelled by user.')
-                return None, None
+            except (EOFError, KeyboardInterrupt, Exception) as e:
+                if isinstance(e, KeyboardInterrupt):
+                    print('\n❌ Operation cancelled by user (Escape/Ctrl+C pressed).')
+                elif isinstance(e, EOFError):
+                    print('\n❌ Operation cancelled by user (EOF detected).')
+                else:
+                    print(f'\n❌ Operation cancelled due to error: {str(e)}')
+                
+                # Explicitly raise SystemExit to ensure the notebook stops
+                raise SystemExit("Infrastructure selection cancelled by user")
 
     def _find_infrastructure_instances(self, infrastructure: INFRASTRUCTURE) -> list[tuple[INFRASTRUCTURE, int | None]]:
         """
@@ -1038,16 +1053,17 @@ def create_resource_group(rg_name: str, resource_group_location: str | None = No
             f"Failed to create the resource group '{rg_name}'", 
             False, False, False, False)
 
-def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int) -> bool:
+def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_update_option: bool = False) -> bool:
     """
     Check if a specific infrastructure exists by querying the resource group.
 
     Args:
         infrastructure (INFRASTRUCTURE): The infrastructure type to check.
         index (int): index for multi-instance infrastructures.
+        allow_update_option (bool): If True, provides option to proceed with infrastructure update when infrastructure exists.
 
     Returns:
-        bool: True if the infrastructure exists, False otherwise.
+        bool: True if the infrastructure exists and no update is desired, False if infrastructure doesn't exist or update is confirmed.
     """
     
     print(f'🔍 Checking if infrastructure already exists...')
@@ -1055,10 +1071,53 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int) -> boo
     rg_name = get_infra_rg_name(infrastructure, index)
 
     if does_resource_group_exist(rg_name):
-        print(f'✅ Infrastructure already exists!\n')
-        print('ℹ️  To redeploy, either:')
-        print('     1. Use a different index number, or')
-        print('     2. Delete the existing resource group first using the clean-up notebook')
+        print(f'✅ Infrastructure already exists: {rg_name}\n')
+        
+        if allow_update_option:
+            print('🔄 Infrastructure Update Options:\n')
+            print('   This infrastructure notebook can update the existing infrastructure.')
+            print('   Updates are additive and will:')
+            print('   • Add new APIs and policy fragments defined in the infrastructure')
+            print('   • Update existing infrastructure components to match the template')
+            print('   • Preserve manually added samples and configurations\n')
+            
+            print('ℹ️  Choose an option:')
+            print('     1. Update the existing infrastructure (recommended - not destructive if samples exist)')
+            print('     2. Use a different index')
+            print('     3. Delete the existing resource group first using the clean-up notebook')
+            
+            while True:
+                try:
+                    choice = input('\nEnter your choice (1, 2, or 3) [default: 1]: ').strip()
+                    
+                    # Default to option 1 if user just presses Enter
+                    if choice == '' or choice == '1':
+                        print('\n🚀 Proceeding with infrastructure update...')
+                        return False  # Allow deployment to proceed
+                    elif choice == '2':
+                        print('\n📝 Please change the index number in the notebook and re-run.')
+                        return True  # Block deployment
+                    elif choice == '3':
+                        print('\n🗑️ Please use the clean-up notebook to delete the resource group first.')
+                        return True  # Block deployment
+                    else:
+                        print('❌ Invalid choice. Please enter 1, 2, or 3 (or press Enter for default).')
+                        
+                except (EOFError, KeyboardInterrupt, Exception) as e:
+                    if isinstance(e, KeyboardInterrupt):
+                        print('\n\n❌ Operation cancelled by user (Escape/Ctrl+C pressed).')
+                    elif isinstance(e, EOFError):
+                        print('\n\n❌ Operation cancelled by user (EOF detected).')
+                    else:
+                        print(f'\n\n❌ Operation cancelled due to error: {str(e)}')
+                    
+                    # Explicitly raise SystemExit to ensure the notebook stops
+                    raise SystemExit("Infrastructure deployment cancelled by user")
+        else:
+            print('ℹ️  To redeploy, either:')
+            print('     1. Use a different index, or')
+            print('     2. Delete the existing resource group first using the clean-up notebook')
+            
         return True
     else:
         print('   Infrastructure does not yet exist.')    

From 930bc1fe24533ebcd440c2a530c036fff2757007 Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 14:37:30 -0400
Subject: [PATCH 2/8] Simplify logic and respect Escape key on inputs

---
 infrastructure/afd-apim-pe/create.ipynb       |   6 +-
 .../afd-apim-pe/create_infrastructure.py      |   5 +-
 infrastructure/apim-aca/create.ipynb          |   6 +-
 .../apim-aca/create_infrastructure.py         |   5 +-
 infrastructure/simple-apim/create.ipynb       |   6 +-
 .../simple-apim/create_infrastructure.py      |   6 +-
 shared/python/infrastructures.py              |  18 +-
 shared/python/utils.py                        | 165 ++++++++++++------
 8 files changed, 142 insertions(+), 75 deletions(-)

diff --git a/infrastructure/afd-apim-pe/create.ipynb b/infrastructure/afd-apim-pe/create.ipynb
index 7fb88b7..20e65da 100644
--- a/infrastructure/afd-apim-pe/create.ipynb
+++ b/infrastructure/afd-apim-pe/create.ipynb
@@ -37,11 +37,7 @@
     "# ------------------------------\n",
     "\n",
     "inb_helper = utils.InfrastructureNotebookHelper(rg_location, INFRASTRUCTURE.AFD_APIM_PE, index, apim_sku) \n",
-    "success = inb_helper.create_infrastructure()\n",
-    "\n",
-    "if not success:\n",
-    "    print(\"❌ Infrastructure creation failed!\")\n",
-    "    raise SystemExit(1)\n",
+    "inb_helper.create_infrastructure()\n",
     "\n",
     "utils.print_ok('All done!')"
    ]
diff --git a/infrastructure/afd-apim-pe/create_infrastructure.py b/infrastructure/afd-apim-pe/create_infrastructure.py
index 2004338..c325375 100644
--- a/infrastructure/afd-apim-pe/create_infrastructure.py
+++ b/infrastructure/afd-apim-pe/create_infrastructure.py
@@ -11,11 +11,14 @@
 
 def create_infrastructure(location: str, index: int, apim_sku: APIM_SKU, no_aca: bool = False) -> None:
     try:
+        # Check if infrastructure already exists to determine messaging
+        infrastructure_exists = utils.does_resource_group_exist(utils.get_infra_rg_name(utils.INFRASTRUCTURE.AFD_APIM_PE, index))
+        
         # Create custom APIs for AFD-APIM-PE with optional Container Apps backends
         custom_apis = _create_afd_specific_apis(not no_aca)
         
         infra = AfdApimAcaInfrastructure(location, index, apim_sku, infra_apis = custom_apis)
-        result = infra.deploy_infrastructure()
+        result = infra.deploy_infrastructure(infrastructure_exists)
         
         sys.exit(0 if result.success else 1)
             
diff --git a/infrastructure/apim-aca/create.ipynb b/infrastructure/apim-aca/create.ipynb
index e4efc7d..7e8b99e 100644
--- a/infrastructure/apim-aca/create.ipynb
+++ b/infrastructure/apim-aca/create.ipynb
@@ -35,11 +35,7 @@
     "# ------------------------------\n",
     "\n",
     "inb_helper = utils.InfrastructureNotebookHelper(rg_location, INFRASTRUCTURE.APIM_ACA, index, apim_sku) \n",
-    "success = inb_helper.create_infrastructure()\n",
-    "\n",
-    "if not success:\n",
-    "    print(\"❌ Infrastructure creation failed!\")\n",
-    "    raise SystemExit(1)\n",
+    "inb_helper.create_infrastructure()\n",
     "\n",
     "utils.print_ok('All done!')"
    ]
diff --git a/infrastructure/apim-aca/create_infrastructure.py b/infrastructure/apim-aca/create_infrastructure.py
index 0c6ea6e..c688f25 100644
--- a/infrastructure/apim-aca/create_infrastructure.py
+++ b/infrastructure/apim-aca/create_infrastructure.py
@@ -11,11 +11,14 @@
 
 def create_infrastructure(location: str, index: int, apim_sku: APIM_SKU) -> None:
     try:
+        # Check if infrastructure already exists to determine messaging
+        infrastructure_exists = utils.does_resource_group_exist(utils.get_infra_rg_name(utils.INFRASTRUCTURE.APIM_ACA, index))
+        
         # Create custom APIs for APIM-ACA with Container Apps backends
         custom_apis = _create_aca_specific_apis()
         
         infra = ApimAcaInfrastructure(location, index, apim_sku, infra_apis = custom_apis)
-        result = infra.deploy_infrastructure()
+        result = infra.deploy_infrastructure(infrastructure_exists)
         
         sys.exit(0 if result.success else 1)
             
diff --git a/infrastructure/simple-apim/create.ipynb b/infrastructure/simple-apim/create.ipynb
index 1c9b7df..cc5f545 100644
--- a/infrastructure/simple-apim/create.ipynb
+++ b/infrastructure/simple-apim/create.ipynb
@@ -35,11 +35,7 @@
     "# ------------------------------\n",
     "\n",
     "inb_helper = utils.InfrastructureNotebookHelper(rg_location, INFRASTRUCTURE.SIMPLE_APIM, index, apim_sku) \n",
-    "success = inb_helper.create_infrastructure()\n",
-    "\n",
-    "if not success:\n",
-    "    print(\"❌ Infrastructure creation failed!\")\n",
-    "    raise SystemExit(1)\n",
+    "inb_helper.create_infrastructure()\n",
     "\n",
     "utils.print_ok('All done!')"
    ]
diff --git a/infrastructure/simple-apim/create_infrastructure.py b/infrastructure/simple-apim/create_infrastructure.py
index 1fc0cd5..f8d29ca 100644
--- a/infrastructure/simple-apim/create_infrastructure.py
+++ b/infrastructure/simple-apim/create_infrastructure.py
@@ -6,11 +6,15 @@
 import argparse
 from apimtypes import APIM_SKU
 from infrastructures import SimpleApimInfrastructure
+import utils
 
 
 def create_infrastructure(location: str, index: int, apim_sku: APIM_SKU) -> None:    
     try:
-        result = SimpleApimInfrastructure(location, index, apim_sku).deploy_infrastructure()
+        # Check if infrastructure already exists to determine messaging
+        infrastructure_exists = utils.does_resource_group_exist(utils.get_infra_rg_name(utils.INFRASTRUCTURE.SIMPLE_APIM, index))
+        
+        result = SimpleApimInfrastructure(location, index, apim_sku).deploy_infrastructure(infrastructure_exists)
         sys.exit(0 if result.success else 1)
             
     except Exception as e:
diff --git a/shared/python/infrastructures.py b/shared/python/infrastructures.py
index b6948da..67250f8 100644
--- a/shared/python/infrastructures.py
+++ b/shared/python/infrastructures.py
@@ -172,13 +172,17 @@ def _verify_infrastructure_specific(self, rg_name: str) -> bool:
     #    PUBLIC METHODS
     # ------------------------------   
 
-    def deploy_infrastructure(self) -> 'utils.Output':
+    def deploy_infrastructure(self, is_update: bool = False) -> 'utils.Output':
         """
         Deploy the infrastructure using the defined Bicep parameters.
         This method should be implemented in subclasses to handle specific deployment logic.
+        
+        Args:
+            is_update (bool): Whether this is an update to existing infrastructure or a new deployment.
         """
         
-        print(f'\n🚀 Creating infrastructure...\n')
+        action_verb = "Updating" if is_update else "Creating"
+        print(f'\n🚀 {action_verb} infrastructure...\n')
         print(f'   Infrastructure : {self.infra.value}')
         print(f'   Index          : {self.index}')
         print(f'   Resource group : {self.rg_name}')
@@ -487,14 +491,18 @@ def _verify_apim_connectivity(self, apim_gateway_url: str) -> bool:
             print('   ℹ️  Continuing deployment - this may be expected during infrastructure setup')
             return True  # Continue anyway
 
-    def deploy_infrastructure(self) -> Output:
+    def deploy_infrastructure(self, is_update: bool = False) -> Output:
         """
         Deploy the AFD-APIM-PE infrastructure with the required multi-step process.
         
+        Args:
+            is_update (bool): Whether this is an update to existing infrastructure or a new deployment.
+            
         Returns:
             utils.Output: The deployment result.
         """
-        print('\n🚀 Starting AFD-APIM-PE infrastructure deployment...\n')
+        action_verb = "Updating" if is_update else "Starting"
+        print(f'\n🚀 {action_verb} AFD-APIM-PE infrastructure deployment...\n')
         print('   This deployment requires multiple steps:\n')
         print('   1. Initial deployment with public access enabled')
         print('   2. Approve private link connections')  
@@ -503,7 +511,7 @@ def deploy_infrastructure(self) -> Output:
         print('   5. Final verification\n')
         
         # Step 1 & 2: Initial deployment using base class method
-        output = super().deploy_infrastructure()
+        output = super().deploy_infrastructure(is_update)
         
         if not output.success:
             print('❌ Initial deployment failed!')
diff --git a/shared/python/utils.py b/shared/python/utils.py
index 889cb55..d584d32 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -266,7 +266,7 @@ def __init__(self, rg_location: str, deployment: INFRASTRUCTURE, index: int, api
     #    PUBLIC METHODS
     # ------------------------------
 
-    def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow_update: bool = True) -> bool:
+    def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow_update: bool = True) -> None:
         """
         Create infrastructure by executing the appropriate creation script.
         
@@ -275,57 +275,80 @@ def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow
             allow_update (bool): Allow infrastructure updates when infrastructure already exists. Defaults to True.
             
         Returns:
-            bool: True if infrastructure creation succeeded, False otherwise.
+            None: Method either succeeds or exits the program with SystemExit.
         """
 
-        import sys 
-
         try:
-            # For infrastructure notebooks, allow the update option
-            infrastructure_exists = does_infrastructure_exist(self.deployment, self.index, allow_update_option=allow_update)
-        except SystemExit:
-            # User cancelled the operation
-            return False
-        
-        if bypass_infrastructure_check or not infrastructure_exists:
-            # Map infrastructure types to their folder names
-            infra_folder_map = {
-                INFRASTRUCTURE.SIMPLE_APIM: 'simple-apim',
-                INFRASTRUCTURE.AFD_APIM_PE: 'afd-apim-pe', 
-                INFRASTRUCTURE.APIM_ACA: 'apim-aca'
-            }
+            import sys 
+
+            # For infrastructure notebooks, check if update is allowed and handle user choice
+            if allow_update:
+                rg_name = get_infra_rg_name(self.deployment, self.index)
+                if does_resource_group_exist(rg_name):
+                    # Infrastructure exists, show update dialog
+                    try:
+                        should_proceed = _prompt_for_infrastructure_update(rg_name)
+                        if not should_proceed:
+                            print('❌ Infrastructure deployment cancelled by user.')
+                            raise SystemExit("User cancelled deployment")
+                    except (KeyboardInterrupt, EOFError):
+                        print('\n❌ Infrastructure deployment cancelled by user (Escape/Ctrl+C pressed).')
+                        # Raise an exception that will stop notebook execution
+                        raise KeyboardInterrupt("User cancelled infrastructure deployment")
             
-            infra_folder = infra_folder_map.get(self.deployment)
-            if not infra_folder:
-                print(f'❌ Unsupported infrastructure type: {self.deployment.value}')
-                return False
+            # Check infrastructure existence for the normal flow
+            infrastructure_exists = does_resource_group_exist(get_infra_rg_name(self.deployment, self.index)) if not allow_update else False
             
-            # Build the command to call the infrastructure creation script
-            cmd_args = [
-                sys.executable, 
-                os.path.join(find_project_root(), 'infrastructure', infra_folder, 'create_infrastructure.py'),
-                '--location', self.rg_location,
-                '--index', str(self.index),
-                '--sku', str(self.apim_sku.value)
-            ]
-
-            # Execute the infrastructure creation script with real-time output streaming and UTF-8 encoding to handle Unicode characters properly
-            process = subprocess.Popen(cmd_args, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, 
-                                    bufsize = 1, universal_newlines = True, encoding = 'utf-8', errors = 'replace')
-
-            try:
-                # Stream output in real-time
-                for line in process.stdout:
-                    print(line.rstrip())
-            except Exception as e:
-                print(f'Error reading subprocess output: {e}')
+            if bypass_infrastructure_check or not infrastructure_exists:
+                # Map infrastructure types to their folder names
+                infra_folder_map = {
+                    INFRASTRUCTURE.SIMPLE_APIM: 'simple-apim',
+                    INFRASTRUCTURE.AFD_APIM_PE: 'afd-apim-pe', 
+                    INFRASTRUCTURE.APIM_ACA: 'apim-aca'
+                }
                 
-            # Wait for process to complete
-            process.wait()
+                infra_folder = infra_folder_map.get(self.deployment)
+                if not infra_folder:
+                    print(f'❌ Unsupported infrastructure type: {self.deployment.value}')
+                    raise SystemExit(1)
+                
+                # Build the command to call the infrastructure creation script
+                cmd_args = [
+                    sys.executable, 
+                    os.path.join(find_project_root(), 'infrastructure', infra_folder, 'create_infrastructure.py'),
+                    '--location', self.rg_location,
+                    '--index', str(self.index),
+                    '--sku', str(self.apim_sku.value)
+                ]
+
+                # Execute the infrastructure creation script with real-time output streaming and UTF-8 encoding to handle Unicode characters properly
+                process = subprocess.Popen(cmd_args, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, 
+                                        bufsize = 1, universal_newlines = True, encoding = 'utf-8', errors = 'replace')
 
-            return process.returncode == 0
+                try:
+                    # Stream output in real-time
+                    for line in process.stdout:
+                        print(line.rstrip())
+                except Exception as e:
+                    print(f'Error reading subprocess output: {e}')
+                    
+                # Wait for process to complete
+                process.wait()
 
-        return True
+                if process.returncode != 0:
+                    print("❌ Infrastructure creation failed!")
+                    raise SystemExit(1)
+
+                return True
+
+            return True
+            
+        except KeyboardInterrupt:
+            print("\n🚫 Infrastructure deployment cancelled by user.")
+            raise SystemExit("User cancelled deployment")
+        except Exception as e:
+            print(f"❌ Infrastructure deployment failed with error: {e}")
+            raise SystemExit(1)
         
 class NotebookHelper:
     """
@@ -1053,6 +1076,45 @@ def create_resource_group(rg_name: str, resource_group_location: str | None = No
             f"Failed to create the resource group '{rg_name}'", 
             False, False, False, False)
 
+def _prompt_for_infrastructure_update(rg_name: str) -> bool:
+    """
+    Prompt the user for infrastructure update confirmation.
+    
+    Args:
+        rg_name (str): The resource group name.
+        
+    Returns:
+        bool: True if user wants to proceed with update, False to cancel.
+    """
+    print(f'✅ Infrastructure already exists: {rg_name}\n')
+    print('🔄 Infrastructure Update Options:\n')
+    print('   This infrastructure notebook can update the existing infrastructure.')
+    print('   Updates are additive and will:')
+    print('   • Add new APIs and policy fragments defined in the infrastructure')
+    print('   • Update existing infrastructure components to match the template')
+    print('   • Preserve manually added samples and configurations\n')
+    
+    print('ℹ️  Choose an option:')
+    print('     1. Update the existing infrastructure (recommended)')
+    print('     2. Use a different index')
+    print('     3. Delete the existing resource group first using the clean-up notebook')
+    
+    while True:
+        choice = input('\nEnter your choice (1, 2, or 3) [default: 1]: ').strip()
+        
+        # Default to option 1 if user just presses Enter
+        if choice == '' or choice == '1':
+            print('\n🚀 Proceeding with infrastructure update...')
+            return True
+        elif choice == '2':
+            print('\n📝 Please change the index number in the notebook and re-run.')
+            return False
+        elif choice == '3':
+            print('\n🗑️ Please use the clean-up notebook to delete the resource group first.')
+            return False
+        else:
+            print('❌ Invalid choice. Please enter 1, 2, or 3 (or press Enter for default).')
+
 def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_update_option: bool = False) -> bool:
     """
     Check if a specific infrastructure exists by querying the resource group.
@@ -1066,7 +1128,8 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
         bool: True if the infrastructure exists and no update is desired, False if infrastructure doesn't exist or update is confirmed.
     """
     
-    print(f'🔍 Checking if infrastructure already exists...')
+    print(f'� Debug: does_infrastructure_exist called with allow_update_option={allow_update_option}')
+    print(f'�🔍 Checking if infrastructure already exists...')
 
     rg_name = get_infra_rg_name(infrastructure, index)
 
@@ -1075,14 +1138,13 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
         
         if allow_update_option:
             print('🔄 Infrastructure Update Options:\n')
-            print('   This infrastructure notebook can update the existing infrastructure.')
-            print('   Updates are additive and will:')
+            print('   This infrastructure notebook can update the existing infrastructure. Updates are additive and will:\n')
             print('   • Add new APIs and policy fragments defined in the infrastructure')
             print('   • Update existing infrastructure components to match the template')
             print('   • Preserve manually added samples and configurations\n')
             
-            print('ℹ️  Choose an option:')
-            print('     1. Update the existing infrastructure (recommended - not destructive if samples exist)')
+            print('ℹ️  Choose an option:\n')
+            print('     1. Update the existing infrastructure (recommended and not destructive if samples already exist)')
             print('     2. Use a different index')
             print('     3. Delete the existing resource group first using the clean-up notebook')
             
@@ -1092,7 +1154,6 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
                     
                     # Default to option 1 if user just presses Enter
                     if choice == '' or choice == '1':
-                        print('\n🚀 Proceeding with infrastructure update...')
                         return False  # Allow deployment to proceed
                     elif choice == '2':
                         print('\n📝 Please change the index number in the notebook and re-run.')
@@ -1111,8 +1172,8 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
                     else:
                         print(f'\n\n❌ Operation cancelled due to error: {str(e)}')
                     
-                    # Explicitly raise SystemExit to ensure the notebook stops
-                    raise SystemExit("Infrastructure deployment cancelled by user")
+                    # Return True to block deployment and let the calling code handle it
+                    return True
         else:
             print('ℹ️  To redeploy, either:')
             print('     1. Use a different index, or')

From ff53cc341cb9de68b51ae86e3347fdef5e15b925 Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 19:26:42 -0400
Subject: [PATCH 3/8] Fix the choice input process

---
 shared/python/utils.py | 69 ++++++++++++------------------------------
 1 file changed, 19 insertions(+), 50 deletions(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index d584d32..4549159 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -293,8 +293,7 @@ def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow
                             raise SystemExit("User cancelled deployment")
                     except (KeyboardInterrupt, EOFError):
                         print('\n❌ Infrastructure deployment cancelled by user (Escape/Ctrl+C pressed).')
-                        # Raise an exception that will stop notebook execution
-                        raise KeyboardInterrupt("User cancelled infrastructure deployment")
+                        raise SystemExit("User cancelled deployment")
             
             # Check infrastructure existence for the normal flow
             infrastructure_exists = does_resource_group_exist(get_infra_rg_name(self.deployment, self.index)) if not allow_update else False
@@ -537,7 +536,7 @@ def _query_and_select_infrastructure(self) -> tuple[INFRASTRUCTURE | None, int |
         # Get user selection
         while True:
             try:
-                choice = input(f'Select infrastructure (1-{len(display_options)}) or press Enter to exit: ').strip()
+                choice = input(f'Select infrastructure (1-{len(display_options)}): ').strip()
                 
                 if not choice:
                     print_warning('No infrastructure selected. Exiting.')
@@ -568,16 +567,7 @@ def _query_and_select_infrastructure(self) -> tuple[INFRASTRUCTURE | None, int |
                     
             except ValueError:
                 print_error('Invalid input. Please enter a number.')
-            except (EOFError, KeyboardInterrupt, Exception) as e:
-                if isinstance(e, KeyboardInterrupt):
-                    print('\n❌ Operation cancelled by user (Escape/Ctrl+C pressed).')
-                elif isinstance(e, EOFError):
-                    print('\n❌ Operation cancelled by user (EOF detected).')
-                else:
-                    print(f'\n❌ Operation cancelled due to error: {str(e)}')
-                
-                # Explicitly raise SystemExit to ensure the notebook stops
-                raise SystemExit("Infrastructure selection cancelled by user")
+
 
     def _find_infrastructure_instances(self, infrastructure: INFRASTRUCTURE) -> list[tuple[INFRASTRUCTURE, int | None]]:
         """
@@ -1097,23 +1087,18 @@ def _prompt_for_infrastructure_update(rg_name: str) -> bool:
     print('ℹ️  Choose an option:')
     print('     1. Update the existing infrastructure (recommended)')
     print('     2. Use a different index')
-    print('     3. Delete the existing resource group first using the clean-up notebook')
+    print('     3. Delete the existing resource group first using the clean-up notebook\n')
     
     while True:
-        choice = input('\nEnter your choice (1, 2, or 3) [default: 1]: ').strip()
+        choice = input('\nEnter your choice (1, 2, or 3): ').strip()
         
         # Default to option 1 if user just presses Enter
-        if choice == '' or choice == '1':
-            print('\n🚀 Proceeding with infrastructure update...')
+        if choice == '1':
             return True
-        elif choice == '2':
-            print('\n📝 Please change the index number in the notebook and re-run.')
-            return False
-        elif choice == '3':
-            print('\n🗑️ Please use the clean-up notebook to delete the resource group first.')
+        elif not choice or choice in('2', '3'):
             return False
         else:
-            print('❌ Invalid choice. Please enter 1, 2, or 3 (or press Enter for default).')
+            print('❌ Invalid choice. Please enter 1, 2, or 3.')
 
 def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_update_option: bool = False) -> bool:
     """
@@ -1146,38 +1131,22 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
             print('ℹ️  Choose an option:\n')
             print('     1. Update the existing infrastructure (recommended and not destructive if samples already exist)')
             print('     2. Use a different index')
-            print('     3. Delete the existing resource group first using the clean-up notebook')
+            print('     3. Delete the existing resource group first using the clean-up notebook\n')
             
             while True:
-                try:
-                    choice = input('\nEnter your choice (1, 2, or 3) [default: 1]: ').strip()
-                    
-                    # Default to option 1 if user just presses Enter
-                    if choice == '' or choice == '1':
-                        return False  # Allow deployment to proceed
-                    elif choice == '2':
-                        print('\n📝 Please change the index number in the notebook and re-run.')
-                        return True  # Block deployment
-                    elif choice == '3':
-                        print('\n🗑️ Please use the clean-up notebook to delete the resource group first.')
-                        return True  # Block deployment
-                    else:
-                        print('❌ Invalid choice. Please enter 1, 2, or 3 (or press Enter for default).')
-                        
-                except (EOFError, KeyboardInterrupt, Exception) as e:
-                    if isinstance(e, KeyboardInterrupt):
-                        print('\n\n❌ Operation cancelled by user (Escape/Ctrl+C pressed).')
-                    elif isinstance(e, EOFError):
-                        print('\n\n❌ Operation cancelled by user (EOF detected).')
-                    else:
-                        print(f'\n\n❌ Operation cancelled due to error: {str(e)}')
-                    
-                    # Return True to block deployment and let the calling code handle it
-                    return True
+                choice = input('\nEnter your choice (1, 2, or 3): ').strip()
+                
+                # Default to option 1 if user just presses Enter
+                if choice == '1':
+                    return False  # Allow deployment to proceed
+                elif not choice or choice in('2', '3'):
+                    return True  # Block deployment
+                else:
+                    print('❌ Invalid choice. Please enter 1, 2, or 3.')
         else:
             print('ℹ️  To redeploy, either:')
             print('     1. Use a different index, or')
-            print('     2. Delete the existing resource group first using the clean-up notebook')
+            print('     2. Delete the existing resource group first using the clean-up notebook\n')
             
         return True
     else:

From 31b94d9df28a4babf8fcc9b140073054970bfaa8 Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 22:32:21 -0400
Subject: [PATCH 4/8] Minor copy change

---
 shared/python/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index 4549159..ccb5f3b 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -1436,7 +1436,7 @@ def cleanup_infra_deployments(deployment: INFRASTRUCTURE, indexes: int | list[in
         print_info(f'Cleaning up resources for {deployment.value} - {idx}', True)
         rg_name = get_infra_rg_name(deployment, idx)
         _cleanup_resources(deployment.value, rg_name)
-        print_ok('Cleanup completed!')
+        print_ok('All done!')
         return
 
     # For multiple indexes, run in parallel

From 329d47e0b8db1c64a7d4163845e47d92302970e4 Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 23:19:30 -0400
Subject: [PATCH 5/8] Retry with a new index when so selected

---
 shared/python/utils.py | 47 ++++++++++++++++++++++++++++++++----------
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index ccb5f3b..c658ef6 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -287,8 +287,14 @@ def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow
                 if does_resource_group_exist(rg_name):
                     # Infrastructure exists, show update dialog
                     try:
-                        should_proceed = _prompt_for_infrastructure_update(rg_name)
-                        if not should_proceed:
+                        should_proceed, new_index = _prompt_for_infrastructure_update(rg_name)
+                        if new_index is not None:
+                            # User selected option 2: Use a different index
+                            print(f'🔄 Retrying infrastructure creation with index {new_index}...\n')
+                            self.index = new_index
+                            # Recursively call create_infrastructure with the new index
+                            return self.create_infrastructure(bypass_infrastructure_check, allow_update)
+                        elif not should_proceed:
                             print('❌ Infrastructure deployment cancelled by user.')
                             raise SystemExit("User cancelled deployment")
                     except (KeyboardInterrupt, EOFError):
@@ -1066,7 +1072,7 @@ def create_resource_group(rg_name: str, resource_group_location: str | None = No
             f"Failed to create the resource group '{rg_name}'", 
             False, False, False, False)
 
-def _prompt_for_infrastructure_update(rg_name: str) -> bool:
+def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
     """
     Prompt the user for infrastructure update confirmation.
     
@@ -1074,7 +1080,9 @@ def _prompt_for_infrastructure_update(rg_name: str) -> bool:
         rg_name (str): The resource group name.
         
     Returns:
-        bool: True if user wants to proceed with update, False to cancel.
+        tuple: (proceed_with_update, new_index) where:
+            - proceed_with_update: True if user wants to proceed with update, False to cancel
+            - new_index: None if no index change, integer if user selected option 2
     """
     print(f'✅ Infrastructure already exists: {rg_name}\n')
     print('🔄 Infrastructure Update Options:\n')
@@ -1087,16 +1095,33 @@ def _prompt_for_infrastructure_update(rg_name: str) -> bool:
     print('ℹ️  Choose an option:')
     print('     1. Update the existing infrastructure (recommended)')
     print('     2. Use a different index')
-    print('     3. Delete the existing resource group first using the clean-up notebook\n')
+    print('     3. Exit, then delete the existing resource group separately via the clean-up notebook\n')
     
     while True:
         choice = input('\nEnter your choice (1, 2, or 3): ').strip()
         
         # Default to option 1 if user just presses Enter
         if choice == '1':
-            return True
-        elif not choice or choice in('2', '3'):
-            return False
+            return True, None
+        elif choice == '2':
+            # Option 2: Prompt for a different index
+            while True:
+                try:
+                    new_index_str = input('\nEnter the desired index for the infrastructure: ').strip()
+                    if not new_index_str:
+                        print('❌ Please enter a valid index number.')
+                        continue
+                    
+                    new_index = int(new_index_str)
+                    if new_index <= 0:
+                        print('❌ Index must be a positive integer.')
+                        continue
+                    
+                    return False, new_index
+                except ValueError:
+                    print('❌ Please enter a valid integer for the index.')
+        elif not choice or choice == '3':
+            return False, None
         else:
             print('❌ Invalid choice. Please enter 1, 2, or 3.')
 
@@ -1131,7 +1156,7 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
             print('ℹ️  Choose an option:\n')
             print('     1. Update the existing infrastructure (recommended and not destructive if samples already exist)')
             print('     2. Use a different index')
-            print('     3. Delete the existing resource group first using the clean-up notebook\n')
+            print('     3. Exit, then delete the existing resource group separately via the clean-up notebook\n')
             
             while True:
                 choice = input('\nEnter your choice (1, 2, or 3): ').strip()
@@ -1146,7 +1171,7 @@ def does_infrastructure_exist(infrastructure: INFRASTRUCTURE, index: int, allow_
         else:
             print('ℹ️  To redeploy, either:')
             print('     1. Use a different index, or')
-            print('     2. Delete the existing resource group first using the clean-up notebook\n')
+            print('     2. Exit, then delete the existing resource group separately via the clean-up notebook\n')
             
         return True
     else:
@@ -1436,7 +1461,7 @@ def cleanup_infra_deployments(deployment: INFRASTRUCTURE, indexes: int | list[in
         print_info(f'Cleaning up resources for {deployment.value} - {idx}', True)
         rg_name = get_infra_rg_name(deployment, idx)
         _cleanup_resources(deployment.value, rg_name)
-        print_ok('All done!')
+        print_ok('Cleanup completed!')
         return
 
     # For multiple indexes, run in parallel

From b20d922dd295cef4619bfeb5aad5735eebe65eef Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Thu, 7 Aug 2025 23:24:47 -0400
Subject: [PATCH 6/8] Allow escaping when selecting a new index

---
 shared/python/utils.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index c658ef6..fcf1ad0 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -290,7 +290,7 @@ def create_infrastructure(self, bypass_infrastructure_check: bool = False, allow
                         should_proceed, new_index = _prompt_for_infrastructure_update(rg_name)
                         if new_index is not None:
                             # User selected option 2: Use a different index
-                            print(f'🔄 Retrying infrastructure creation with index {new_index}...\n')
+                            print(f'🔄 Retrying infrastructure creation with index {new_index}...')
                             self.index = new_index
                             # Recursively call create_infrastructure with the new index
                             return self.create_infrastructure(bypass_infrastructure_check, allow_update)
@@ -1086,13 +1086,12 @@ def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
     """
     print(f'✅ Infrastructure already exists: {rg_name}\n')
     print('🔄 Infrastructure Update Options:\n')
-    print('   This infrastructure notebook can update the existing infrastructure.')
-    print('   Updates are additive and will:')
+    print('   This infrastructure notebook can update the existing infrastructure. Updates are additive and will:\n')
     print('   • Add new APIs and policy fragments defined in the infrastructure')
     print('   • Update existing infrastructure components to match the template')
     print('   • Preserve manually added samples and configurations\n')
     
-    print('ℹ️  Choose an option:')
+    print('ℹ️  Choose an option:\n')
     print('     1. Update the existing infrastructure (recommended)')
     print('     2. Use a different index')
     print('     3. Exit, then delete the existing resource group separately via the clean-up notebook\n')
@@ -1109,8 +1108,7 @@ def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
                 try:
                     new_index_str = input('\nEnter the desired index for the infrastructure: ').strip()
                     if not new_index_str:
-                        print('❌ Please enter a valid index number.')
-                        continue
+                        return False, None
                     
                     new_index = int(new_index_str)
                     if new_index <= 0:

From 8916fd7fe946911cdfe95dafa9bae3febaa92f43 Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Fri, 8 Aug 2025 00:23:52 -0400
Subject: [PATCH 7/8] Fix warning

---
 tests/python/pytest.ini | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/python/pytest.ini b/tests/python/pytest.ini
index 5293117..8708905 100644
--- a/tests/python/pytest.ini
+++ b/tests/python/pytest.ini
@@ -7,7 +7,6 @@ markers =
     slow: tests that take a long time to run
     unit: marks tests as unit tests
     http: marks tests that mock or use HTTP
-testpaths =
-    tests/python
+testpaths = .
 python_files =
     test_*.py

From a3f135198fa3fcedc10aa3df3ff5112d251ee89c Mon Sep 17 00:00:00 2001
From: Simon Kurtz <simonkurtz@microsoft.com>
Date: Fri, 8 Aug 2025 00:24:07 -0400
Subject: [PATCH 8/8] Format, copy changes

---
 shared/python/utils.py     |  15 +--
 tests/python/test_utils.py | 232 +++++++++++++++++++++++++++++++++----
 2 files changed, 219 insertions(+), 28 deletions(-)

diff --git a/shared/python/utils.py b/shared/python/utils.py
index fcf1ad0..b8c04ad 100644
--- a/shared/python/utils.py
+++ b/shared/python/utils.py
@@ -1086,21 +1086,22 @@ def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
     """
     print(f'✅ Infrastructure already exists: {rg_name}\n')
     print('🔄 Infrastructure Update Options:\n')
-    print('   This infrastructure notebook can update the existing infrastructure. Updates are additive and will:\n')
+    print('   This infrastructure notebook can update the existing infrastructure.')
+    print('   Updates are additive and will:')
     print('   • Add new APIs and policy fragments defined in the infrastructure')
     print('   • Update existing infrastructure components to match the template')
     print('   • Preserve manually added samples and configurations\n')
     
-    print('ℹ️  Choose an option:\n')
+    print('ℹ️  Choose an option:')
     print('     1. Update the existing infrastructure (recommended)')
     print('     2. Use a different index')
-    print('     3. Exit, then delete the existing resource group separately via the clean-up notebook\n')
+    print('     3. Delete the existing resource group first using the clean-up notebook\n')
     
     while True:
         choice = input('\nEnter your choice (1, 2, or 3): ').strip()
         
         # Default to option 1 if user just presses Enter
-        if choice == '1':
+        if choice == '1' or not choice:
             return True, None
         elif choice == '2':
             # Option 2: Prompt for a different index
@@ -1108,7 +1109,8 @@ def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
                 try:
                     new_index_str = input('\nEnter the desired index for the infrastructure: ').strip()
                     if not new_index_str:
-                        return False, None
+                        print('❌ Please enter a valid index number.')
+                        continue
                     
                     new_index = int(new_index_str)
                     if new_index <= 0:
@@ -1118,7 +1120,7 @@ def _prompt_for_infrastructure_update(rg_name: str) -> tuple[bool, int | None]:
                     return False, new_index
                 except ValueError:
                     print('❌ Please enter a valid integer for the index.')
-        elif not choice or choice == '3':
+        elif choice == '3':
             return False, None
         else:
             print('❌ Invalid choice. Please enter 1, 2, or 3.')
@@ -1523,7 +1525,6 @@ def cleanup_infra_deployments(deployment: INFRASTRUCTURE, indexes: int | list[in
                     print_error(f"❌ Exception during cleanup for {deployment.value}-{task['index']}: {str(e)}")
 
     # Final summary
-    print()
     if failed_count == 0:
         print_ok(f'All {len(indexes_list)} infrastructure cleanups completed successfully!')
     else:
diff --git a/tests/python/test_utils.py b/tests/python/test_utils.py
index 44a7194..1cf44a3 100644
--- a/tests/python/test_utils.py
+++ b/tests/python/test_utils.py
@@ -1686,11 +1686,19 @@ def test_query_and_select_infrastructure_no_options(monkeypatch):
     monkeypatch.setattr(nb_helper, '_find_infrastructure_instances', lambda x: [])
     monkeypatch.setattr(utils, 'print_info', lambda *args, **kwargs: None)
     monkeypatch.setattr(utils, 'print_warning', lambda *args, **kwargs: None)
-    # Mock input to return empty string (simulating user pressing Enter to exit)
-    monkeypatch.setattr('builtins.input', lambda prompt: '')
+    monkeypatch.setattr(utils, 'print_success', lambda *args, **kwargs: None)
+    
+    # Mock the infrastructure creation to succeed
+    def mock_infrastructure_creation(self, bypass_check=True):
+        return True
+    
+    monkeypatch.setattr(utils.InfrastructureNotebookHelper, 'create_infrastructure', mock_infrastructure_creation)
     
+    # When no infrastructures are available, it should automatically create new infrastructure
     result = nb_helper._query_and_select_infrastructure()
-    assert result == (None, None)
+    
+    # Expect it to return the desired infrastructure and None index (since 'test-rg' doesn't match the expected pattern)
+    assert result == (INFRASTRUCTURE.SIMPLE_APIM, None)
 
 def test_query_and_select_infrastructure_single_option(monkeypatch):
     """Test _query_and_select_infrastructure with a single available option."""
@@ -1818,30 +1826,212 @@ def mock_infrastructure_creation(self, bypass_check=True):
     result = nb_helper._query_and_select_infrastructure()
     assert result == (INFRASTRUCTURE.SIMPLE_APIM, 2)
 
-def test_query_and_select_infrastructure_keyboard_interrupt(monkeypatch):
-    """Test _query_and_select_infrastructure when user presses Ctrl+C."""
-    nb_helper = utils.NotebookHelper(
-        'test-sample', 'test-rg', 'eastus', 
-        INFRASTRUCTURE.SIMPLE_APIM, [INFRASTRUCTURE.SIMPLE_APIM]
-    )
+
+# ------------------------------
+#    TESTS FOR _prompt_for_infrastructure_update
+# ------------------------------
+
+def test_prompt_for_infrastructure_update_option_1(monkeypatch):
+    """Test _prompt_for_infrastructure_update when user selects option 1 (update)."""
+    monkeypatch.setattr('builtins.input', lambda prompt: '1')
     
-    # Mock single result
-    def mock_find_instances(infra):
-        return [(INFRASTRUCTURE.SIMPLE_APIM, 1)]
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (True, None)
+
+def test_prompt_for_infrastructure_update_option_1_default(monkeypatch):
+    """Test _prompt_for_infrastructure_update when user presses Enter (defaults to option 1)."""
+    monkeypatch.setattr('builtins.input', lambda prompt: '')
     
-    monkeypatch.setattr(nb_helper, '_find_infrastructure_instances', mock_find_instances)
-    monkeypatch.setattr(utils, 'print_info', lambda *args, **kwargs: None)
-    monkeypatch.setattr(utils, 'print_warning', lambda *args, **kwargs: None)
-    monkeypatch.setattr(utils, 'get_infra_rg_name', lambda infra, idx: f'apim-infra-{infra.value}-{idx}')
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (True, None)
+
+def test_prompt_for_infrastructure_update_option_2_valid_index(monkeypatch):
+    """Test _prompt_for_infrastructure_update when user selects option 2 with valid index."""
+    inputs = iter(['2', '5'])  # Option 2, then index 5
+    monkeypatch.setattr('builtins.input', lambda prompt: next(inputs))
+    
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (False, 5)
+
+def test_prompt_for_infrastructure_update_option_2_invalid_then_valid_index(monkeypatch):
+    """Test _prompt_for_infrastructure_update when user provides invalid index then valid one."""
+    inputs = iter(['2', '', '0', '-1', 'abc', '3'])  # Option 2, then empty, zero, negative, non-number, finally valid
+    monkeypatch.setattr('builtins.input', lambda prompt: next(inputs))
+    
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (False, 3)
+
+def test_prompt_for_infrastructure_update_option_3(monkeypatch):
+    """Test _prompt_for_infrastructure_update when user selects option 3 (delete first)."""
+    monkeypatch.setattr('builtins.input', lambda prompt: '3')
+    
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (False, None)
+
+def test_prompt_for_infrastructure_update_invalid_choice_then_valid(monkeypatch):
+    """Test _prompt_for_infrastructure_update with invalid choice followed by valid choice."""
+    inputs = iter(['4', '0', 'invalid', '1'])  # Invalid choices, then option 1
+    monkeypatch.setattr('builtins.input', lambda prompt: next(inputs))
+    
+    result = utils._prompt_for_infrastructure_update('test-rg')
+    assert result == (True, None)
+
+
+# ------------------------------
+#    TESTS FOR InfrastructureNotebookHelper.create_infrastructure WITH INDEX RETRY
+# ------------------------------
+
+def test_infrastructure_notebook_helper_create_with_index_retry(monkeypatch):
+    """Test InfrastructureNotebookHelper.create_infrastructure with option 2 (different index) retry."""
+    from apimtypes import INFRASTRUCTURE, APIM_SKU
+    
+    helper = utils.InfrastructureNotebookHelper('eastus', INFRASTRUCTURE.SIMPLE_APIM, 1, APIM_SKU.BASICV2)
+    
+    # Mock resource group existence to return True initially
+    call_count = 0
+    def mock_rg_exists(rg_name):
+        nonlocal call_count
+        call_count += 1
+        # First call (index 1) returns True, second call (index 3) returns False
+        return call_count == 1
+    
+    # Mock the prompt to return option 2 with index 3
+    monkeypatch.setattr(utils, '_prompt_for_infrastructure_update', lambda rg_name: (False, 3))
+    monkeypatch.setattr(utils, 'does_resource_group_exist', mock_rg_exists)
+    
+    # Mock subprocess execution to succeed
+    class MockProcess:
+        def __init__(self, *args, **kwargs):
+            self.returncode = 0
+            self.stdout = iter(['Mock deployment output\n', 'Success!\n'])
+        
+        def wait(self):
+            pass
+    
+    monkeypatch.setattr('subprocess.Popen', MockProcess)
+    monkeypatch.setattr(utils, 'find_project_root', lambda: 'c:\\mock\\root')
+    
+    # Mock print functions to avoid output during testing
+    monkeypatch.setattr('builtins.print', lambda *args, **kwargs: None)
+    
+    # Should succeed after retrying with index 3
+    result = helper.create_infrastructure()
+    assert result is True
+    assert helper.index == 3  # Verify index was updated
+
+def test_infrastructure_notebook_helper_create_with_recursive_retry(monkeypatch):
+    """Test InfrastructureNotebookHelper.create_infrastructure with multiple recursive retries."""
+    from apimtypes import INFRASTRUCTURE, APIM_SKU
+    
+    helper = utils.InfrastructureNotebookHelper('eastus', INFRASTRUCTURE.SIMPLE_APIM, 1, APIM_SKU.BASICV2)
+    
+    # Mock resource group existence for multiple indexes
+    rg_checks = {}
+    def mock_rg_exists(rg_name):
+        # Parse index from resource group name
+        if 'simple-apim-1' in rg_name:
+            return True  # Index 1 exists
+        elif 'simple-apim-2' in rg_name:
+            return True  # Index 2 also exists
+        else:
+            return False  # Index 3 doesn't exist
+    
+    # Mock the prompt to first return index 2, then index 3
+    prompt_calls = 0
+    def mock_prompt(rg_name):
+        nonlocal prompt_calls
+        prompt_calls += 1
+        if prompt_calls == 1:
+            return (False, 2)  # First retry with index 2
+        else:
+            return (False, 3)  # Second retry with index 3
+    
+    monkeypatch.setattr(utils, '_prompt_for_infrastructure_update', mock_prompt)
+    monkeypatch.setattr(utils, 'does_resource_group_exist', mock_rg_exists)
+    
+    # Mock subprocess execution to succeed
+    class MockProcess:
+        def __init__(self, *args, **kwargs):
+            self.returncode = 0
+            self.stdout = iter(['Mock deployment output\n'])
+        
+        def wait(self):
+            pass
+    
+    monkeypatch.setattr('subprocess.Popen', MockProcess)
+    monkeypatch.setattr(utils, 'find_project_root', lambda: 'c:\\mock\\root')
+    monkeypatch.setattr('builtins.print', lambda *args, **kwargs: None)
+    
+    # Should succeed after retrying with index 3
+    result = helper.create_infrastructure()
+    assert result is True
+    assert helper.index == 3  # Verify final index
+
+def test_infrastructure_notebook_helper_create_user_cancellation(monkeypatch):
+    """Test InfrastructureNotebookHelper.create_infrastructure when user cancels during retry."""
+    from apimtypes import INFRASTRUCTURE, APIM_SKU
+    import pytest
+    
+    helper = utils.InfrastructureNotebookHelper('eastus', INFRASTRUCTURE.SIMPLE_APIM, 1, APIM_SKU.BASICV2)
+    
+    # Mock resource group to exist (triggering prompt)
+    monkeypatch.setattr(utils, 'does_resource_group_exist', lambda rg_name: True)
+    
+    # Mock the prompt to return cancellation (option 3)
+    monkeypatch.setattr(utils, '_prompt_for_infrastructure_update', lambda rg_name: (False, None))
     monkeypatch.setattr('builtins.print', lambda *args, **kwargs: None)
     
-    # Mock user input to raise KeyboardInterrupt
-    def mock_input(prompt):
+    # Should raise SystemExit when user cancels
+    with pytest.raises(SystemExit) as exc_info:
+        helper.create_infrastructure()
+    
+    assert "User cancelled deployment" in str(exc_info.value)
+
+def test_infrastructure_notebook_helper_create_keyboard_interrupt_during_prompt(monkeypatch):
+    """Test InfrastructureNotebookHelper.create_infrastructure when KeyboardInterrupt occurs during prompt."""
+    from apimtypes import INFRASTRUCTURE, APIM_SKU
+    import pytest
+    
+    helper = utils.InfrastructureNotebookHelper('eastus', INFRASTRUCTURE.SIMPLE_APIM, 1, APIM_SKU.BASICV2)
+    
+    # Mock resource group to exist (triggering prompt)
+    monkeypatch.setattr(utils, 'does_resource_group_exist', lambda rg_name: True)
+    
+    # Mock the prompt to raise KeyboardInterrupt
+    def mock_prompt(rg_name):
         raise KeyboardInterrupt()
-    monkeypatch.setattr('builtins.input', mock_input)
     
-    result = nb_helper._query_and_select_infrastructure()
-    assert result == (None, None)
+    monkeypatch.setattr(utils, '_prompt_for_infrastructure_update', mock_prompt)
+    monkeypatch.setattr('builtins.print', lambda *args, **kwargs: None)
+    
+    # Should raise SystemExit when KeyboardInterrupt occurs
+    with pytest.raises(SystemExit) as exc_info:
+        helper.create_infrastructure()
+    
+    assert "User cancelled deployment" in str(exc_info.value)
+
+def test_infrastructure_notebook_helper_create_eof_error_during_prompt(monkeypatch):
+    """Test InfrastructureNotebookHelper.create_infrastructure when EOFError occurs during prompt."""
+    from apimtypes import INFRASTRUCTURE, APIM_SKU
+    import pytest
+    
+    helper = utils.InfrastructureNotebookHelper('eastus', INFRASTRUCTURE.SIMPLE_APIM, 1, APIM_SKU.BASICV2)
+    
+    # Mock resource group to exist (triggering prompt)
+    monkeypatch.setattr(utils, 'does_resource_group_exist', lambda rg_name: True)
+    
+    # Mock the prompt to raise EOFError
+    def mock_prompt(rg_name):
+        raise EOFError()
+    
+    monkeypatch.setattr(utils, '_prompt_for_infrastructure_update', mock_prompt)
+    monkeypatch.setattr('builtins.print', lambda *args, **kwargs: None)
+    
+    # Should raise SystemExit when EOFError occurs
+    with pytest.raises(SystemExit) as exc_info:
+        helper.create_infrastructure()
+    
+    assert "User cancelled deployment" in str(exc_info.value)
 
 def test_deploy_sample_with_infrastructure_selection(monkeypatch):
     """Test deploy_sample method with infrastructure selection when original doesn't exist."""