From ddcd474a5e2ce4568cca646eb1f5bce32b4ba0ed Mon Sep 17 00:00:00 2001
From: Shahar Epstein <60007259+shahar1@users.noreply.github.com>
Date: Sun, 30 Jul 2023 07:57:22 +0300
Subject: [PATCH] Fix check if `virtualenv` is installed in
 `PythonVirtualenvOperator` (#32939)

---
 airflow/operators/python.py    |  4 ++--
 tests/operators/test_python.py | 10 ++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/airflow/operators/python.py b/airflow/operators/python.py
index 574fe5b0e34f..3a10cf4f7362 100644
--- a/airflow/operators/python.py
+++ b/airflow/operators/python.py
@@ -17,11 +17,11 @@
 # under the License.
 from __future__ import annotations
 
+import importlib
 import inspect
 import logging
 import os
 import pickle
-import shutil
 import subprocess
 import sys
 import types
@@ -540,7 +540,7 @@ def __init__(
                 "major versions for PythonVirtualenvOperator. Please use string_args."
                 f"Sys version: {sys.version_info}. Venv version: {python_version}"
             )
-        if not shutil.which("virtualenv"):
+        if importlib.util.find_spec("virtualenv") is None:
             raise AirflowException("PythonVirtualenvOperator requires virtualenv, please install it.")
         if not requirements:
             self.requirements: list[str] | str = []
diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py
index 534a6be5a992..ba296f77a405 100644
--- a/tests/operators/test_python.py
+++ b/tests/operators/test_python.py
@@ -847,6 +847,16 @@ def default_kwargs(*, python_version=sys.version_info[0], **kwargs):
         kwargs["python_version"] = python_version
         return kwargs
 
+    @mock.patch("airflow.operators.python.importlib")
+    def test_virtuenv_not_installed(self, importlib):
+        importlib.util.find_spec.return_value = None
+        with pytest.raises(AirflowException, match="requires virtualenv"):
+
+            def f():
+                pass
+
+            self.run_as_task(f)
+
     def test_add_dill(self):
         def f():
             """Ensure dill is correctly installed."""