modified test_env_getter

tests/integration/connections/debug_envgetter.py

@@ -2,15 +2,13 @@
 import json
 import logging
 
+from databricks.labs.remorph.connections.env_getter import EnvGetter
 
-class EnvGetter:
-    def __init__(self, is_debug: bool = True):
-        self.env = self._get_debug_env() if is_debug else dict(os.environ)
 
-    def get(self, key: str) -> str:
-        if key in self.env:
-            return self.env[key]
-        raise KeyError(f"not in env: {key}")
+class TestEnvGetter(EnvGetter):
+    def __init__(self, is_debug: bool = True):
+        self.is_debug = is_debug
+        super().__init__()
 
     def _get_debug_env(self) -> dict:
         try:
@@ -22,3 +20,8 @@ def _get_debug_env(self) -> dict:
             return raw.get("ucws", {})
         except FileNotFoundError:
             return dict(os.environ)
+
+    def get(self, key: str) -> str:
+        if self.is_debug:
+            self.env = self._get_debug_env()
+        return super().get(key)

tests/integration/connections/test_mssql_connector.py

@@ -5,7 +5,7 @@
 
 from databricks.labs.remorph.connections.credential_manager import create_credential_manager
 from databricks.labs.remorph.connections.database_manager import DatabaseManager, MSSQLConnector
-from .debug_envgetter import EnvGetter
+from .debug_envgetter import TestEnvGetter
 
 
 @pytest.fixture(scope="module")
@@ -29,7 +29,8 @@ def mock_credentials():
 
 @pytest.fixture(scope="module")
 def db_manager(mock_credentials):
-    config = create_credential_manager("remorph", EnvGetter(True)).get_credentials("mssql")
+    env = TestEnvGetter(True)
+    config = create_credential_manager("remorph", env).get_credentials("mssql")
     # since the kv has only URL so added explicit parse rules
     base_url, params = config['server'].replace("jdbc:", "", 1).split(";", 1)