diff --git a/tests/utils/test_decorators.py b/tests/utils/test_decorators.py index 19d3ec31d0311..a0a8ea263427b 100644 --- a/tests/utils/test_decorators.py +++ b/tests/utils/test_decorators.py @@ -49,7 +49,7 @@ def test_task_decorator_using_source(decorator: TaskDecorator): def f(): return ["some_task"] - assert parse_python_source(f, "decorator") == 'def f():\n return ["some_task"]\n' + assert parse_python_source(f, "decorator") == "def f():\n return ['some_task']" @pytest.mark.parametrize("decorator", DECORATORS, indirect=["decorator"]) @@ -59,7 +59,7 @@ def test_skip_if(decorator: TaskDecorator): def f(): return "hello world" - assert parse_python_source(f, "decorator") == 'def f():\n return "hello world"\n' + assert parse_python_source(f, "decorator") == "def f():\n return 'hello world'" @pytest.mark.parametrize("decorator", DECORATORS, indirect=["decorator"]) @@ -69,7 +69,7 @@ def test_run_if(decorator: TaskDecorator): def f(): return "hello world" - assert parse_python_source(f, "decorator") == 'def f():\n return "hello world"\n' + assert parse_python_source(f, "decorator") == "def f():\n return 'hello world'" def test_skip_if_and_run_if(): @@ -79,7 +79,7 @@ def test_skip_if_and_run_if(): def f(): return "hello world" - assert parse_python_source(f) == 'def f():\n return "hello world"\n' + assert parse_python_source(f) == "def f():\n return 'hello world'" def test_run_if_and_skip_if(): @@ -89,7 +89,7 @@ def test_run_if_and_skip_if(): def f(): return "hello world" - assert parse_python_source(f) == 'def f():\n return "hello world"\n' + assert parse_python_source(f) == "def f():\n return 'hello world'" def test_skip_if_allow_decorator(): @@ -102,7 +102,7 @@ def non_task_decorator(func): def f(): return "hello world" - assert parse_python_source(f) == '@non_task_decorator\ndef f():\n return "hello world"\n' + assert parse_python_source(f) == "@non_task_decorator\ndef f():\n return 'hello world'" def test_run_if_allow_decorator(): @@ -115,7 +115,7 @@ def non_task_decorator(func): def f(): return "hello world" - assert parse_python_source(f) == '@non_task_decorator\ndef f():\n return "hello world"\n' + assert parse_python_source(f) == "@non_task_decorator\ndef f():\n return 'hello world'" def parse_python_source(task: Task, custom_operator_name: str | None = None) -> str: diff --git a/tests/utils/test_preexisting_python_virtualenv_decorator.py b/tests/utils/test_preexisting_python_virtualenv_decorator.py index 11d80e348ea81..1a8aa0fa8229b 100644 --- a/tests/utils/test_preexisting_python_virtualenv_decorator.py +++ b/tests/utils/test_preexisting_python_virtualenv_decorator.py @@ -22,20 +22,20 @@ class TestExternalPythonDecorator: def test_remove_task_decorator(self): - py_source = '@task.external_python(serializer="dill")\ndef f():\nimport funcsigs' + py_source = "@task.external_python(use_dill=True)\ndef f(): ...\nimport funcsigs" res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.external_python") - assert res == "def f():\nimport funcsigs" + assert res == "def f():\n ...\nimport funcsigs" def test_remove_decorator_no_parens(self): - py_source = "@task.external_python\ndef f():\nimport funcsigs" + py_source = "@task.external_python\ndef f(): ...\nimport funcsigs" res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.external_python") - assert res == "def f():\nimport funcsigs" + assert res == "def f():\n ...\nimport funcsigs" def test_remove_decorator_nested(self): - py_source = "@foo\n@task.external_python\n@bar\ndef f():\nimport funcsigs" + py_source = "@foo\n@task.external_python\n@bar\ndef f(): ...\nimport funcsigs" res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.external_python") - assert res == "@foo\n@bar\ndef f():\nimport funcsigs" + assert res == "@foo\n@bar\ndef f():\n ...\nimport funcsigs" - py_source = "@foo\n@task.external_python()\n@bar\ndef f():\nimport funcsigs" + py_source = "@foo\n@task.external_python()\n@bar\ndef f(): ...\nimport funcsigs" res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.external_python") - assert res == "@foo\n@bar\ndef f():\nimport funcsigs" + assert res == "@foo\n@bar\ndef f():\n ...\nimport funcsigs" diff --git a/tests/utils/test_python_virtualenv.py b/tests/utils/test_python_virtualenv.py new file mode 100644 index 0000000000000..4d683b7ddd1f2 --- /dev/null +++ b/tests/utils/test_python_virtualenv.py @@ -0,0 +1,135 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import sys +from pathlib import Path +from unittest import mock + +import pytest + +from airflow.utils.decorators import remove_task_decorator +from airflow.utils.python_virtualenv import _generate_pip_conf, prepare_virtualenv + + +class TestPrepareVirtualenv: + @pytest.mark.parametrize( + ("index_urls", "expected_pip_conf_content", "unexpected_pip_conf_content"), + [ + [[], ["[global]", "no-index ="], ["index-url", "extra", "http", "pypi"]], + [["http://mysite"], ["[global]", "index-url", "http://mysite"], ["no-index", "extra", "pypi"]], + [ + ["http://mysite", "https://othersite"], + ["[global]", "index-url", "http://mysite", "extra", "https://othersite"], + ["no-index", "pypi"], + ], + [ + ["http://mysite", "https://othersite", "http://site"], + ["[global]", "index-url", "http://mysite", "extra", "https://othersite http://site"], + ["no-index", "pypi"], + ], + ], + ) + def test_generate_pip_conf( + self, + index_urls: list[str], + expected_pip_conf_content: list[str], + unexpected_pip_conf_content: list[str], + tmp_path: Path, + ): + tmp_file = tmp_path / "pip.conf" + _generate_pip_conf(tmp_file, index_urls) + generated_conf = tmp_file.read_text() + for term in expected_pip_conf_content: + assert term in generated_conf + for term in unexpected_pip_conf_content: + assert term not in generated_conf + + @mock.patch("airflow.utils.python_virtualenv.execute_in_subprocess") + def test_should_create_virtualenv(self, mock_execute_in_subprocess): + python_bin = prepare_virtualenv( + venv_directory="/VENV", python_bin="pythonVER", system_site_packages=False, requirements=[] + ) + assert "/VENV/bin/python" == python_bin + mock_execute_in_subprocess.assert_called_once_with( + [sys.executable, "-m", "virtualenv", "/VENV", "--python=pythonVER"] + ) + + @mock.patch("airflow.utils.python_virtualenv.execute_in_subprocess") + def test_should_create_virtualenv_with_system_packages(self, mock_execute_in_subprocess): + python_bin = prepare_virtualenv( + venv_directory="/VENV", python_bin="pythonVER", system_site_packages=True, requirements=[] + ) + assert "/VENV/bin/python" == python_bin + mock_execute_in_subprocess.assert_called_once_with( + [sys.executable, "-m", "virtualenv", "/VENV", "--system-site-packages", "--python=pythonVER"] + ) + + @mock.patch("airflow.utils.python_virtualenv.execute_in_subprocess") + def test_pip_install_options(self, mock_execute_in_subprocess): + pip_install_options = ["--no-deps"] + python_bin = prepare_virtualenv( + venv_directory="/VENV", + python_bin="pythonVER", + system_site_packages=True, + requirements=["apache-beam[gcp]"], + pip_install_options=pip_install_options, + ) + + assert "/VENV/bin/python" == python_bin + mock_execute_in_subprocess.assert_any_call( + [sys.executable, "-m", "virtualenv", "/VENV", "--system-site-packages", "--python=pythonVER"] + ) + mock_execute_in_subprocess.assert_called_with( + ["/VENV/bin/pip", "install", *pip_install_options, "apache-beam[gcp]"] + ) + + @mock.patch("airflow.utils.python_virtualenv.execute_in_subprocess") + def test_should_create_virtualenv_with_extra_packages(self, mock_execute_in_subprocess): + python_bin = prepare_virtualenv( + venv_directory="/VENV", + python_bin="pythonVER", + system_site_packages=False, + requirements=["apache-beam[gcp]"], + ) + assert "/VENV/bin/python" == python_bin + + mock_execute_in_subprocess.assert_any_call( + [sys.executable, "-m", "virtualenv", "/VENV", "--python=pythonVER"] + ) + + mock_execute_in_subprocess.assert_called_with(["/VENV/bin/pip", "install", "apache-beam[gcp]"]) + + def test_remove_task_decorator(self): + py_source = "@task.virtualenv(use_dill=True)\ndef f(): ...\nimport funcsigs" + res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv") + assert res == "def f():\n ...\nimport funcsigs" + + def test_remove_decorator_no_parens(self): + py_source = "@task.virtualenv\ndef f(): ...\nimport funcsigs" + res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv") + assert res == "def f():\n ...\nimport funcsigs" + + def test_remove_decorator_nested(self): + py_source = "@foo\n@task.virtualenv\n@bar\ndef f(): ...\nimport funcsigs" + res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv") + assert res == "@foo\n@bar\ndef f():\n ...\nimport funcsigs" + + py_source = "@foo\n@task.virtualenv()\n@bar\ndef f(): ...\nimport funcsigs" + res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv") + assert res == "@foo\n@bar\ndef f():\n ...\nimport funcsigs"