josix commented on code in PR #43383:
URL: https://github.com/apache/airflow/pull/43383#discussion_r1820101887


##########
tests/utils/test_python_virtualenv.py:
##########
@@ -116,25 +116,20 @@ def 
test_should_create_virtualenv_with_extra_packages(self, mock_execute_in_subp
         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"
+        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():\nimport funcsigs"
+        assert res == "def f():\n    ...\nimport funcsigs"
 
     def test_remove_decorator_no_parens(self):
-        py_source = "@task.virtualenv\ndef f():\nimport funcsigs"
+        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():\nimport funcsigs"
-
-    def test_remove_decorator_including_comment(self):
-        py_source = "@task.virtualenv\ndef f():\n# @task.virtualenv\nimport 
funcsigs"
-        res = remove_task_decorator(python_source=py_source, 
task_decorator_name="@task.virtualenv")
-        assert res == "def f():\n# @task.virtualenv\nimport funcsigs"

Review Comment:
   Yeah, I think it's better to leave dag authors' code information as much as 
possible in the rendered `script.py` of virtualenv operator, ideally the unit 
tests should not be changed, let me try LibCST instead of AST here.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to