chore: Add class tear down for cleaning up custom credentials (#4805)

* Add class tear down for cleaning up custom credentials

* Fix failing teardownclass and refactor setup

* Address mypy suggestion

* Refactoring the teardown to be a test case tear down

* Fix mypy issue

---------

Co-authored-by: Wing Fung Lau <[email protected]>

tests/integration/init/schemas/schemas_test_data_setup.py

@@ -18,13 +18,23 @@
 
 
 class SchemaTestDataSetup(TestCase):
-    original_cred_file = None
-    original_config_file = None
-    original_profile = None
-    original_region = None
+    original_cred_file: str
+    original_config_file: str
+    original_profile: str
+    original_region: str
 
     @classmethod
     def setUpClass(cls):
+        env = os.environ
+        if AWS_CONFIG_FILE in env:
+            cls.original_config_file = env[AWS_CONFIG_FILE]
+        if AWS_SHARED_CREDENTIALS_FILE in env:
+            cls.original_cred_file = env[AWS_SHARED_CREDENTIALS_FILE]
+        if AWS_PROFILE in env:
+            cls.original_profile = env[AWS_PROFILE]
+        if AWS_DEFAULT_REGION in env:
+            cls.original_region = env[AWS_DEFAULT_REGION]
+
         session = Session()
         schemas_client = session.client("schemas", region_name=session.region_name)
         # all setup is done here to avoid creating side effects in test. Currently we are using CLI and
@@ -55,17 +65,33 @@ def setUpClass(cls):
             runner = CliRunner()
             runner.invoke(init_cmd, ["--output-dir", temp], input=user_input)
 
+    def tearDown(self) -> None:
+        env = os.environ
+        if env.get(AWS_CONFIG_FILE):
+            del env[AWS_CONFIG_FILE]
+        if self.original_config_file:
+            env[AWS_CONFIG_FILE] = self.original_config_file
+
+        if env.get(AWS_SHARED_CREDENTIALS_FILE):
+            del env[AWS_SHARED_CREDENTIALS_FILE]
+        if self.original_cred_file:
+            env[AWS_SHARED_CREDENTIALS_FILE] = self.original_cred_file
+
+        if env.get(AWS_PROFILE):
+            del env[AWS_PROFILE]
+        if self.original_profile:
+            env[AWS_PROFILE] = self.original_profile
+
+        if env.get(AWS_DEFAULT_REGION):
+            del env[AWS_DEFAULT_REGION]
+        if self.original_region:
+            env[AWS_DEFAULT_REGION] = self.original_region
+
+        shutil.rmtree(self.config_dir, ignore_errors=True)
+
     def _init_custom_config(self, profile, region):
         self.config_dir = tempfile.mkdtemp()
         env = os.environ
-        if AWS_CONFIG_FILE in env:
-            self.original_config_file = env[AWS_CONFIG_FILE]
-        if AWS_SHARED_CREDENTIALS_FILE in env:
-            self.original_cred_file = env[AWS_SHARED_CREDENTIALS_FILE]
-        if AWS_PROFILE in env:
-            self.original_profile = env[AWS_PROFILE]
-        if AWS_DEFAULT_REGION in env:
-            self.original_region = env[AWS_DEFAULT_REGION]
 
         custom_config = self._create_config_file(profile, region)
         session = Session()
@@ -81,31 +107,6 @@ def _init_custom_config(self, profile, region):
         env[AWS_PROFILE] = profile
         env[AWS_DEFAULT_REGION] = region
 
-    def _tear_down_custom_config(self):
-        env = os.environ
-
-        if self.original_config_file is None:
-            del env[AWS_CONFIG_FILE]
-        else:
-            env[AWS_CONFIG_FILE] = self.original_config_file
-
-        if self.original_cred_file is None:
-            del env[AWS_SHARED_CREDENTIALS_FILE]
-        else:
-            env[AWS_SHARED_CREDENTIALS_FILE] = self.original_cred_file
-
-        if self.original_profile is None:
-            del env[AWS_PROFILE]
-        else:
-            env[AWS_PROFILE] = self.original_profile
-
-        if self.original_region is None:
-            del env[AWS_DEFAULT_REGION]
-        else:
-            env[AWS_DEFAULT_REGION] = self.original_region
-
-        shutil.rmtree(self.config_dir, ignore_errors=True)
-
     def _create_config_file(self, profile, region):
         if profile == DEFAULT:
             config_file_content = "[{0}]\noutput = json\nregion = {1}".format(profile, region)

tests/integration/init/schemas/test_init_with_schemas_command.py

@@ -308,8 +308,6 @@ def test_init_interactive_with_event_bridge_app_non_default_profile_selection(se
             self.assertTrue(expected_output_folder.is_dir())
             self.assertTrue(Path(expected_output_folder, "hello_world_function", "schema").is_dir())
 
-            self._tear_down_custom_config()
-
     def test_init_interactive_with_event_bridge_app_non_supported_schemas_region(self):
         self._init_custom_config("default", "cn-north-1")
         # WHEN the user follows interactive init prompts
@@ -339,4 +337,3 @@ def test_init_interactive_with_event_bridge_app_non_supported_schemas_region(sel
             runner = CliRunner()
             result = runner.invoke(init_cmd, ["--output-dir", temp], input=user_input)
             self.assertTrue(result.exception)
-            self._tear_down_custom_config()