|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
| 15 | +import inspect |
| 16 | +from unittest.mock import patch |
| 17 | + |
15 | 18 | import bigframes_vendored.constants as constants |
16 | 19 | import pytest |
17 | 20 |
|
18 | 21 | from bigframes.functions import _utils, function_typing |
19 | 22 |
|
20 | 23 |
|
| 24 | +def test_get_updated_package_requirements_no_extra_package(): |
| 25 | + """Tests with no extra package.""" |
| 26 | + result = _utils.get_updated_package_requirements(capture_references=False) |
| 27 | + assert result is None |
| 28 | + |
| 29 | + initial_packages = ["xgboost"] |
| 30 | + result = _utils.get_updated_package_requirements( |
| 31 | + initial_packages, capture_references=False |
| 32 | + ) |
| 33 | + assert result == initial_packages |
| 34 | + |
| 35 | + |
| 36 | +@patch("bigframes.functions._utils.numpy.__version__", "1.24.4") |
| 37 | +@patch("bigframes.functions._utils.pyarrow.__version__", "14.0.1") |
| 38 | +@patch("bigframes.functions._utils.pandas.__version__", "2.0.3") |
| 39 | +@patch("bigframes.functions._utils.cloudpickle.__version__", "2.2.1") |
| 40 | +def test_get_updated_package_requirements_is_row_processor_with_versions(): |
| 41 | + """Tests with is_row_processor=True and specific versions.""" |
| 42 | + expected = [ |
| 43 | + "cloudpickle==2.2.1", |
| 44 | + "numpy==1.24.4", |
| 45 | + "pandas==2.0.3", |
| 46 | + "pyarrow==14.0.1", |
| 47 | + ] |
| 48 | + result = _utils.get_updated_package_requirements(is_row_processor=True) |
| 49 | + assert result == expected |
| 50 | + |
| 51 | + |
| 52 | +@patch("bigframes.functions._utils.warnings.warn") |
| 53 | +@patch("bigframes.functions._utils.cloudpickle.__version__", "2.2.1") |
| 54 | +def test_get_updated_package_requirements_ignore_version(mock_warn): |
| 55 | + """ |
| 56 | + Tests with is_row_processor=True and ignore_package_version=True. |
| 57 | + Should add packages without versions and raise a warning. |
| 58 | + """ |
| 59 | + expected = ["cloudpickle==2.2.1", "numpy", "pandas", "pyarrow"] |
| 60 | + result = _utils.get_updated_package_requirements( |
| 61 | + is_row_processor=True, ignore_package_version=True |
| 62 | + ) |
| 63 | + assert result == expected |
| 64 | + # Verify that a warning was issued. |
| 65 | + mock_warn.assert_called_once() |
| 66 | + |
| 67 | + |
| 68 | +@patch("bigframes.functions._utils.numpy.__version__", "1.24.4") |
| 69 | +@patch("bigframes.functions._utils.pyarrow.__version__", "14.0.1") |
| 70 | +@patch("bigframes.functions._utils.pandas.__version__", "2.0.3") |
| 71 | +def test_get_updated_package_requirements_capture_references_false(): |
| 72 | + """ |
| 73 | + Tests with capture_references=False. |
| 74 | + Should not add cloudpickle but should add others if requested. |
| 75 | + """ |
| 76 | + # Case 1: Only capture_references=False. |
| 77 | + result_1 = _utils.get_updated_package_requirements(capture_references=False) |
| 78 | + assert result_1 is None |
| 79 | + |
| 80 | + # Case 2: capture_references=False but is_row_processor=True. |
| 81 | + expected_2 = ["numpy==1.24.4", "pandas==2.0.3", "pyarrow==14.0.1"] |
| 82 | + result_2 = _utils.get_updated_package_requirements( |
| 83 | + is_row_processor=True, capture_references=False |
| 84 | + ) |
| 85 | + assert result_2 == expected_2 |
| 86 | + |
| 87 | + |
| 88 | +@patch("bigframes.functions._utils.numpy.__version__", "1.24.4") |
| 89 | +@patch("bigframes.functions._utils.pyarrow.__version__", "14.0.1") |
| 90 | +@patch("bigframes.functions._utils.pandas.__version__", "2.0.3") |
| 91 | +@patch("bigframes.functions._utils.cloudpickle.__version__", "2.2.1") |
| 92 | +def test_get_updated_package_requirements_non_overlapping_packages(): |
| 93 | + """Tests providing an initial list of packages that do not overlap.""" |
| 94 | + initial_packages = ["scikit-learn==1.3.0", "xgboost"] |
| 95 | + expected = [ |
| 96 | + "cloudpickle==2.2.1", |
| 97 | + "numpy==1.24.4", |
| 98 | + "pandas==2.0.3", |
| 99 | + "pyarrow==14.0.1", |
| 100 | + "scikit-learn==1.3.0", |
| 101 | + "xgboost", |
| 102 | + ] |
| 103 | + result = _utils.get_updated_package_requirements( |
| 104 | + package_requirements=initial_packages, is_row_processor=True |
| 105 | + ) |
| 106 | + assert result == expected |
| 107 | + |
| 108 | + |
| 109 | +@patch("bigframes.functions._utils.numpy.__version__", "1.24.4") |
| 110 | +@patch("bigframes.functions._utils.pyarrow.__version__", "14.0.1") |
| 111 | +@patch("bigframes.functions._utils.pandas.__version__", "2.0.3") |
| 112 | +@patch("bigframes.functions._utils.cloudpickle.__version__", "2.2.1") |
| 113 | +def test_get_updated_package_requirements_overlapping_packages(): |
| 114 | + """Tests that packages are not added if they already exist.""" |
| 115 | + # The function should respect the pre-existing pandas version. |
| 116 | + initial_packages = ["pandas==1.5.3", "numpy"] |
| 117 | + expected = [ |
| 118 | + "cloudpickle==2.2.1", |
| 119 | + "numpy", |
| 120 | + "pandas==1.5.3", |
| 121 | + "pyarrow==14.0.1", |
| 122 | + ] |
| 123 | + result = _utils.get_updated_package_requirements( |
| 124 | + package_requirements=initial_packages, is_row_processor=True |
| 125 | + ) |
| 126 | + assert result == expected |
| 127 | + |
| 128 | + |
| 129 | +@patch("bigframes.functions._utils.cloudpickle.__version__", "2.2.1") |
| 130 | +def test_get_updated_package_requirements_with_existing_cloudpickle(): |
| 131 | + """Tests that cloudpickle is not added if it already exists.""" |
| 132 | + initial_packages = ["cloudpickle==2.0.0"] |
| 133 | + expected = ["cloudpickle==2.0.0"] |
| 134 | + result = _utils.get_updated_package_requirements( |
| 135 | + package_requirements=initial_packages |
| 136 | + ) |
| 137 | + assert result == expected |
| 138 | + |
| 139 | + |
| 140 | +def test_package_existed_helper(): |
| 141 | + """Tests the _package_existed helper function directly.""" |
| 142 | + reqs = ["pandas==1.0", "numpy", "scikit-learn>=1.2.0"] |
| 143 | + # Exact match |
| 144 | + assert _utils._package_existed(reqs, "pandas==1.0") |
| 145 | + # Different version |
| 146 | + assert _utils._package_existed(reqs, "pandas==2.0") |
| 147 | + # No version specified |
| 148 | + assert _utils._package_existed(reqs, "numpy") |
| 149 | + # Not in list |
| 150 | + assert not _utils._package_existed(reqs, "xgboost") |
| 151 | + # Empty list |
| 152 | + assert not _utils._package_existed([], "pandas") |
| 153 | + |
| 154 | + |
| 155 | +def test_has_conflict_output_type_no_conflict(): |
| 156 | + """Tests has_conflict_output_type with type annotation.""" |
| 157 | + # Helper functions with type annotation for has_conflict_output_type. |
| 158 | + def _func_with_return_type(x: int) -> int: |
| 159 | + return x |
| 160 | + |
| 161 | + signature = inspect.signature(_func_with_return_type) |
| 162 | + |
| 163 | + assert _utils.has_conflict_output_type(signature, output_type=float) |
| 164 | + assert not _utils.has_conflict_output_type(signature, output_type=int) |
| 165 | + |
| 166 | + |
| 167 | +def test_has_conflict_output_type_no_annotation(): |
| 168 | + """Tests has_conflict_output_type without type annotation.""" |
| 169 | + # Helper functions without type annotation for has_conflict_output_type. |
| 170 | + def _func_without_return_type(x): |
| 171 | + return x |
| 172 | + |
| 173 | + signature = inspect.signature(_func_without_return_type) |
| 174 | + |
| 175 | + assert not _utils.has_conflict_output_type(signature, output_type=int) |
| 176 | + assert not _utils.has_conflict_output_type(signature, output_type=float) |
| 177 | + |
| 178 | + |
21 | 179 | @pytest.mark.parametrize( |
22 | 180 | ["metadata_options", "metadata_string"], |
23 | 181 | ( |
|
0 commit comments