From cf99e697a87f3fe132ed51c7f9c2476eb0f24d38 Mon Sep 17 00:00:00 2001
From: Lili Deng <lildeng@microsoft.com>
Date: Sat, 14 Sep 2024 13:20:12 +0800
Subject: [PATCH] add unit test test_env_retry_on_deployment_fail

---
 selftests/azure/test_prepare.py       |  3 +-
 selftests/runners/test_lisa_runner.py | 45 +++++++++++++++++++++++++++
 selftests/test_environment.py         |  1 +
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/selftests/azure/test_prepare.py b/selftests/azure/test_prepare.py
index 92e171c31..dcb11e5f9 100644
--- a/selftests/azure/test_prepare.py
+++ b/selftests/azure/test_prepare.py
@@ -368,6 +368,7 @@ class AzurePrepareTestCase(TestCase):
     def load_environment(
         self,
         node_req_count: int = 2,
+        retry: int = 0,
     ) -> Environment:
         runbook = schema.Environment()
         if node_req_count > 0:
@@ -381,7 +382,7 @@ class AzurePrepareTestCase(TestCase):
             warn_as_error=False,
             id_=0,
             runbook=runbook,
-            retry=0,
+            retry=retry,
         )
 
         return environment
diff --git a/selftests/runners/test_lisa_runner.py b/selftests/runners/test_lisa_runner.py
index 640d00240..3492edce5 100644
--- a/selftests/runners/test_lisa_runner.py
+++ b/selftests/runners/test_lisa_runner.py
@@ -616,6 +616,51 @@ class RunnerTestCase(TestCase):
             test_results=test_results,
         )
 
+    def test_env_retry_on_deployment_fail(self) -> None:
+        # retry when fail to deploy
+        platform_schema = test_platform.MockPlatformSchema(deploy_success=False)
+        test_testsuite.generate_cases_metadata()
+        env_runbook = generate_env_runbook()
+        env_runbook.retry = 2
+        runner = generate_runner(env_runbook, platform_schema=platform_schema)
+        test_result_messages = self._run_all_tests(runner)
+        deployment_failure = "deployment failed. LisaException: mock deploy failed"
+        self.verify_env_results(
+            expected_prepared=[
+                "generated_0",
+                "generated_1",
+                "generated_2",
+                "generated_0",
+                "generated_0",
+                "generated_1",
+                "generated_1",
+                "generated_2",
+                "generated_2",
+            ],
+            expected_deployed_envs=[],
+            expected_deleted_envs=[
+                "generated_0",
+                "generated_1",
+                "generated_2",
+            ],
+            runner=runner,
+        )
+        self.verify_test_results(
+            expected_test_order=["mock_ut1", "mock_ut2", "mock_ut3"],
+            expected_envs=["generated_0", "generated_1", "generated_2"],
+            expected_status=[
+                TestStatus.FAILED,
+                TestStatus.FAILED,
+                TestStatus.FAILED,
+            ],
+            expected_message=[
+                deployment_failure,
+                deployment_failure,
+                deployment_failure,
+            ],
+            test_results=test_result_messages,
+        )
+
     def verify_test_results(
         self,
         expected_test_order: List[str],
diff --git a/selftests/test_environment.py b/selftests/test_environment.py
index 46590d8f1..ef8691e6d 100644
--- a/selftests/test_environment.py
+++ b/selftests/test_environment.py
@@ -198,6 +198,7 @@ class EnvironmentTestCase(TestCase):
         envs = load_environments(None)
         self.assertEqual(0, len(envs))
         self.assertEqual(False, envs.warn_as_error)
+        self.assertEqual(0, envs.retry)
 
     def test_create_from_runbook_split(self) -> None:
         runbook = generate_runbook(local=True, remote=True)