Update test_catalog_notebook.py, test_catalog_project.py, and test_vault_vault.py

Author: qiuosier

tests/unitary/default_setup/catalog/test_catalog_notebook.py

@@ -17,10 +17,11 @@
 from oci.exceptions import ServiceError
 
 import ads.config
+import ads.catalog.notebook
 from ads.catalog.notebook import NotebookCatalog, NotebookSummaryList
 from ads.common import auth, oci_client
 from ads.common.utils import random_valid_ocid
-from ads.config import NB_SESSION_COMPARTMENT_OCID, PROJECT_OCID
+from ads.config import PROJECT_OCID
 
 
 def generate_notebook_list(
@@ -82,17 +83,21 @@ class NotebookCatalogTest(unittest.TestCase):
 
             nsl = NotebookSummaryList(generate_notebook_list())
 
-    def setUp(self) -> None:
+    @classmethod
+    def setUpClass(cls) -> None:
         os.environ[
             "NB_SESSION_COMPARTMENT_OCID"
         ] = "ocid1.compartment.oc1.<unique_ocid>"
         reload(ads.config)
-        return super().setUp()
+        reload(ads.catalog.notebook)
+        return super().setUpClass()
 
-    def tearDown(self) -> None:
+    @classmethod
+    def tearDownClass(cls) -> None:
         os.environ.pop("NB_SESSION_COMPARTMENT_OCID", None)
         reload(ads.config)
-        return super().tearDown()
+        reload(ads.catalog.notebook)
+        return super().tearDownClass()
 
     @staticmethod
     def generate_notebook_response_data(compartment_id=None, notebook_id=None):
@@ -124,7 +129,10 @@ def test_notebook_init_with_compartment_id(self, mock_client, mock_signer):
     def test_notebook_init_without_compartment_id(self, mock_client, mock_signer):
         """Test notebook catalog initiation without compartment_id."""
         test_notebook_catalog = NotebookCatalog()
-        assert test_notebook_catalog.compartment_id == NB_SESSION_COMPARTMENT_OCID
+        assert (
+            test_notebook_catalog.compartment_id
+            == ads.config.NB_SESSION_COMPARTMENT_OCID
+        )
 
     def test_decorate_notebook_session_attributes(self):
         """Test NotebookCatalog._decorate_notebook_session method."""

tests/unitary/default_setup/catalog/test_catalog_project.py

@@ -15,10 +15,10 @@
 from oci.exceptions import ServiceError
 
 import ads.config
+import ads.catalog.project
 from ads.catalog.project import ProjectCatalog, ProjectSummaryList
 from ads.common import auth, oci_client
 from ads.common.utils import random_valid_ocid
-from ads.config import NB_SESSION_COMPARTMENT_OCID
 
 
 def generate_project_list(
@@ -76,17 +76,21 @@ class ProjectCatalogTest(unittest.TestCase):
 
             psl = ProjectSummaryList(generate_project_list())
 
-    def setUp(self) -> None:
+    @classmethod
+    def setUpClass(cls) -> None:
         os.environ[
             "NB_SESSION_COMPARTMENT_OCID"
         ] = "ocid1.compartment.oc1.<unique_ocid>"
         reload(ads.config)
-        return super().setUp()
+        reload(ads.catalog.project)
+        return super().setUpClass()
 
-    def tearDown(self) -> None:
+    @classmethod
+    def tearDownClass(cls) -> None:
         os.environ.pop("NB_SESSION_COMPARTMENT_OCID", None)
         reload(ads.config)
-        return super().tearDown()
+        reload(ads.catalog.project)
+        return super().tearDownClass()
 
     @staticmethod
     def generate_project_response_data(compartment_id=None, project_id=None):
@@ -116,7 +120,10 @@ def test_project_init_with_compartment_id(self, mock_client, mock_signer):
     def test_project_init_without_compartment_id(self, mock_client, mock_signer):
         """Test project catalog initiation without compartment_id."""
         test_project_catalog = ProjectCatalog()
-        assert test_project_catalog.compartment_id == NB_SESSION_COMPARTMENT_OCID
+        assert (
+            test_project_catalog.compartment_id
+            == ads.config.NB_SESSION_COMPARTMENT_OCID
+        )
 
     def test_decorate_project_session_attributes(self):
         """Test ProjectCatalog._decorate_project method."""

tests/unitary/with_extras/vault/test_vault_vault.py

@@ -14,19 +14,13 @@
 from oci.vault.models import Secret
 
 import ads.config
+import ads.vault.vault
 from ads.vault.vault import Vault
 
 
 class TestVault:
     """Contains test cases for ads.vault.vault.py"""
 
-    vault = Vault(
-        vault_id="ocid1.vault.oc1.iad.<unique_id>",
-        key_id="ocid1.key.oc1.iad.<unique_id>",
-    )
-    vault.secret_client = MagicMock()
-    vault.vaults_client_composite = MagicMock()
-
     credential = {
         "database_name": "db201910031555_high",
         "username": "admin",
@@ -42,15 +36,25 @@ class TestVault:
     secret_ocid = "ocid1.vaultsecret.oc1.iad.<unique_id>"
     date_time = datetime(2021, 7, 13, 18, 24, 42, 110000, tzinfo=timezone.utc)
 
-    def setup_method(self):
+    @classmethod
+    def setup_class(cls):
         os.environ[
             "NB_SESSION_COMPARTMENT_OCID"
         ] = "ocid1.compartment.oc1.<unique_ocid>"
         reload(ads.config)
+        reload(ads.vault.vault)
+        cls.vault = Vault(
+            vault_id="ocid1.vault.oc1.iad.<unique_id>",
+            key_id="ocid1.key.oc1.iad.<unique_id>",
+        )
+        cls.vault.secret_client = MagicMock()
+        cls.vault.vaults_client_composite = MagicMock()
 
-    def teardown_method(self):
+    @classmethod
+    def teardown_class(cls):
         os.environ.pop("NB_SESSION_COMPARTMENT_OCID", None)
         reload(ads.config)
+        reload(ads.vault.vault)
 
     def test_create_secret(self):
         """Test vault.create_secret()."""