fixup! refactor(utils/decorators): rewrite remove task decorator to u…

…se ast

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:

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"

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"