From bab0e7b266b58a1e743a81ecb4300b037019e29d Mon Sep 17 00:00:00 2001
From: smokestacklightnin
 <125844868+smokestacklightnin@users.noreply.github.com>
Date: Fri, 13 Dec 2024 20:29:44 -0800
Subject: [PATCH] Add mock input for removing stage files

---
 tests/tests_unit/test_upgrade.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/tests_unit/test_upgrade.py b/tests/tests_unit/test_upgrade.py
index 9505b8d554..fe7079f750 100644
--- a/tests/tests_unit/test_upgrade.py
+++ b/tests/tests_unit/test_upgrade.py
@@ -50,6 +50,8 @@ def test_upgrade_4_0(
     monkeypatch,
 ):
     def mock_input(prompt, **kwargs):
+        from _nebari.upgrade import TERRAFORM_REMOVE_TERRAFORM_STAGE_FILES_CONFIRMATION
+
         # Mock different upgrade steps prompt answers
         if prompt == "Have you deleted the Argo Workflows CRDs and service accounts?":
             return True
@@ -68,6 +70,8 @@ def mock_input(prompt, **kwargs):
             == "[bold]Would you like Nebari to assign the corresponding role to all of your current groups automatically?[/bold]"
         ):
             return False
+        elif prompt == TERRAFORM_REMOVE_TERRAFORM_STAGE_FILES_CONFIRMATION:
+            return attempt_fixes
         # All other prompts will be answered with "y"
         else:
             return True