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_python_virtualenv.py b/tests/utils/test_python_virtualenv.py
index 38cda4854baf6..4d683b7ddd1f2 100644
--- a/tests/utils/test_python_virtualenv.py
+++ b/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"
+        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"
+        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():\nimport funcsigs"
+        assert res == "@foo\n@bar\ndef f():\n    ...\nimport funcsigs"
 
-        py_source = "@foo\n@task.virtualenv()\n@bar\ndef f():\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():\nimport funcsigs"
+        assert res == "@foo\n@bar\ndef f():\n    ...\nimport funcsigs"