diff --git a/tests/unit/functions/test_remote_function_utils.py b/tests/unit/functions/test_remote_function_utils.py index 8ddd39d857..812d65bbad 100644 --- a/tests/unit/functions/test_remote_function_utils.py +++ b/tests/unit/functions/test_remote_function_utils.py @@ -13,6 +13,7 @@ # limitations under the License. import inspect +import sys from unittest.mock import patch import bigframes_vendored.constants as constants @@ -227,6 +228,26 @@ def test_get_updated_package_requirements_with_existing_cloudpickle(): assert result == expected +# Dynamically generate expected python versions for the test +_major = sys.version_info.major +_minor = sys.version_info.minor +_compat_version = f"python{_major}{_minor}" +_standard_version = f"python-{_major}.{_minor}" + + +@pytest.mark.parametrize( + "is_compat, expected_version", + [ + (True, _compat_version), + (False, _standard_version), + ], +) +def test_get_python_version(is_compat, expected_version): + """Tests the python version for both standard and compat modes.""" + result = _utils.get_python_version(is_compat=is_compat) + assert result == expected_version + + def test_package_existed_helper(): """Tests the _package_existed helper function directly.""" reqs = ["pandas==1.0", "numpy", "scikit-learn>=1.2.0"]