diff --git a/tests/test_fireboltdialect.py b/tests/test_fireboltdialect.py index a613ccc..7cf7f5c 100644 --- a/tests/test_fireboltdialect.py +++ b/tests/test_fireboltdialect.py @@ -9,6 +9,8 @@ from sqlalchemy import create_engine from sqlalchemy.dialects import registry +from unittest import mock + test_username = os.environ["username"] test_password = os.environ["password"] @@ -28,19 +30,19 @@ def get_engine(): class TestFireboltDialect: def test_create_connect_args(self): - os.environ["FIREBOLT_BASE_URL"] = "test_url" connection_url = "test_engine://test_user@email:test_password@test_db_name/test_engine_name" u = url.make_url(connection_url) - result_list, result_dict = dialect.create_connect_args(u) - assert result_dict["engine_name"] == "test_engine_name" - assert result_dict["username"] == "test_user@email" - assert result_dict["password"] == "test_password" - assert result_dict["database"] == "test_db_name" - assert result_dict["api_endpoint"] == "test_url" + with mock.patch.dict(os.environ, {"FIREBOLT_BASE_URL": "test_url"}): + result_list, result_dict = dialect.create_connect_args(u) + assert result_dict["engine_name"] == "test_engine_name" + assert result_dict["username"] == "test_user@email" + assert result_dict["password"] == "test_password" + assert result_dict["database"] == "test_db_name" + assert result_dict["api_endpoint"] == "test_url" # No endpoint override - del os.environ["FIREBOLT_BASE_URL"] - result_list, result_dict = dialect.create_connect_args(u) - assert "api_endpoint" not in result_dict + with mock.patch.dict(os.environ, {}, clear=True): + result_list, result_dict = dialect.create_connect_args(u) + assert "api_endpoint" not in result_dict def test_get_schema_names(self, get_engine): engine = get_engine