diff --git a/tests/operators/test_virtualenv.py b/tests/operators/test_virtualenv.py index 142a251a7..67ac977bd 100644 --- a/tests/operators/test_virtualenv.py +++ b/tests/operators/test_virtualenv.py @@ -65,3 +65,28 @@ def test_run_command( assert dbt_deps[0][0][0] == dbt_cmd[0][0][0] assert dbt_cmd[0][0][1] == "do-something" assert mock_execute.call_count == 2 + +@patch("airflow.hooks.base.BaseHook.get_connection") +def test_supply_virtualenv_dir_flag( + mock_get_connection, +): + mock_get_connection.return_value = Connection( + conn_id="fake_conn", + conn_type="postgres", + host="fake_host", + port=5432, + login="fake_login", + password="fake_password", + schema="fake_schema", + ) + venv_operator = DbtVirtualenvBaseOperator( + profile_config=profile_config, + task_id="fake_task", + install_deps=True, + project_dir="./dev/dags/dbt/jaffle_shop", + py_system_site_packages=False, + py_requirements=["dbt-postgres==1.6.0b1"], + emit_datasets=False, + virtualenv_dir="mock-venv", + ) + assert venv_operator.venv_dbt_path == "mock-venv/bin/dbt" \ No newline at end of file