diff --git a/cosmos/operators/virtualenv.py b/cosmos/operators/virtualenv.py index b9ae9dd64..97a63f45b 100644 --- a/cosmos/operators/virtualenv.py +++ b/cosmos/operators/virtualenv.py @@ -105,7 +105,7 @@ def _get_or_create_venv_py_interpreter(self) -> str: 1. No: Create one in this directory and return it 2. Yes: Return this virtual env 2. No: Create a temporary virtual env and return it - + """ if self._venv_dir is not None: if self._venv_dir.is_dir(): diff --git a/dev/dags/example_virtualenv.py b/dev/dags/example_virtualenv.py index a46473c3c..45af6e0a8 100644 --- a/dev/dags/example_virtualenv.py +++ b/dev/dags/example_virtualenv.py @@ -25,17 +25,18 @@ ), ) + @dag( schedule_interval="@daily", start_date=datetime(2023, 1, 1), catchup=False, ) def example_virtualenv() -> None: - start_task = EmptyOperator(task_id='start-venv-examples') - end_task = EmptyOperator(task_id='end-venv-examples') + start_task = EmptyOperator(task_id="start-venv-examples") + end_task = EmptyOperator(task_id="end-venv-examples") tmp_venv_task_group = DbtTaskGroup( - group_id='tmp-venv-group', + group_id="tmp-venv-group", # dbt/cosmos-specific parameters project_config=ProjectConfig( DBT_ROOT_PATH / "jaffle_shop", @@ -55,7 +56,7 @@ def example_virtualenv() -> None: ) cached_venv_task_group = DbtTaskGroup( - group_id='cached-venv-group', + group_id="cached-venv-group", # dbt/cosmos-specific parameters project_config=ProjectConfig( DBT_ROOT_PATH / "jaffle_shop", @@ -76,6 +77,6 @@ def example_virtualenv() -> None: start_task >> [tmp_venv_task_group, cached_venv_task_group] >> end_task + example_virtualenv() # [END virtualenv_example] - diff --git a/tests/test_converter.py b/tests/test_converter.py index 5e81aaccf..d50ea1649 100644 --- a/tests/test_converter.py +++ b/tests/test_converter.py @@ -70,6 +70,7 @@ def test_converter_creates_dag_with_seed(mock_load_dbt_graph, execution_mode, op ) assert converter + @pytest.mark.parametrize( "execution_mode,virtualenv_dir,operator_args", [ @@ -81,7 +82,7 @@ def test_converter_creates_dag_with_seed(mock_load_dbt_graph, execution_mode, op @patch("cosmos.converter.DbtGraph.load") def test_converter_raises_warning(mock_load_dbt_graph, execution_mode, virtualenv_dir, operator_args, caplog): """ - This test will raise a warning if we are trying to pass ExecutionMode != `VirtualEnv` + This test will raise a warning if we are trying to pass ExecutionMode != `VirtualEnv` and still pass a defined `virtualenv_dir` """ project_config = ProjectConfig(dbt_project_path=SAMPLE_DBT_PROJECT) @@ -92,7 +93,7 @@ def test_converter_raises_warning(mock_load_dbt_graph, execution_mode, virtualen target_name="my_target_name", profiles_yml_filepath=SAMPLE_PROFILE_YML, ) - converter = DbtToAirflowConverter( + DbtToAirflowConverter( nodes=nodes, project_config=project_config, profile_config=profile_config, @@ -100,6 +101,9 @@ def test_converter_raises_warning(mock_load_dbt_graph, execution_mode, virtualen render_config=render_config, operator_args=operator_args, ) - - assert "`ExecutionConfig.virtualenv_dir` is only supported when \ - ExecutionConfig.execution_mode is set to ExecutionMode.VIRTUALENV." in caplog.text \ No newline at end of file + + assert ( + "`ExecutionConfig.virtualenv_dir` is only supported when \ + ExecutionConfig.execution_mode is set to ExecutionMode.VIRTUALENV." + in caplog.text + )