Update test_vault_vault.py

Author: qiuosier

tests/unitary/with_extras/vault/test_vault_vault.py

@@ -2,15 +2,19 @@
 
 # Copyright (c) 2021, 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
-
-from ads.vault.vault import Vault
+import base64
+import json
+import os
 from collections import namedtuple
 from datetime import datetime, timezone
+from importlib import reload
+from unittest.mock import MagicMock
+
 from oci.secrets.models import SecretBundle, Base64SecretBundleContentDetails
 from oci.vault.models import Secret
-from unittest.mock import MagicMock
-import base64
-import json
+
+import ads.config
+from ads.vault.vault import Vault
 
 
 class TestVault:
@@ -38,6 +42,16 @@ 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):
+        os.environ[
+            "NB_SESSION_COMPARTMENT_OCID"
+        ] = "ocid1.compartment.oc1.<unique_ocid>"
+        reload(ads.config)
+
+    def teardown_method(self):
+        os.environ.pop("NB_SESSION_COMPARTMENT_OCID", None)
+        reload(ads.config)
+
     def test_create_secret(self):
         """Test vault.create_secret()."""
         wrapper = namedtuple("wrapper", ["data"])