Skip to content

Commit f04a626

Browse files
committed
refactor: Consolidate all _utils unit tests of bigframes function
1 parent 46994d7 commit f04a626

File tree

2 files changed

+158
-175
lines changed

2 files changed

+158
-175
lines changed

tests/unit/functions/test_remote_function_utils.py

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,170 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import inspect
16+
from unittest.mock import patch
17+
1518
import bigframes_vendored.constants as constants
1619
import pytest
1720

1821
from bigframes.functions import _utils, function_typing
1922

2023

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+
21179
@pytest.mark.parametrize(
22180
["metadata_options", "metadata_string"],
23181
(

tests/unit/functions/test_utils.py

Lines changed: 0 additions & 175 deletions
This file was deleted.

0 commit comments

Comments
 (0)