Skip to content

Commit 320a02b

Browse files
committed
updates tests for coverage reasons
1 parent c8101e0 commit 320a02b

File tree

3 files changed

+102
-48
lines changed

3 files changed

+102
-48
lines changed

db_dtypes/__init__.py

Lines changed: 18 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,14 @@
2121
import warnings
2222

2323
import numpy
24-
import packaging.version
2524
import pandas
2625
import pandas.api.extensions
2726
from pandas.errors import OutOfBoundsDatetime
2827
import pyarrow
2928
import pyarrow.compute
3029

3130
from db_dtypes import core
32-
from db_dtypes.version import __version__
31+
from db_dtypes.json import JSONArray, JSONDtype, JSONArrowType # noqa: F401
3332

3433
from . import _versions_helpers
3534

@@ -47,15 +46,6 @@
4746
_NP_BOX_DTYPE = "datetime64[us]"
4847

4948

50-
# To use JSONArray and JSONDtype, you'll need Pandas 1.5.0 or later. With the removal
51-
# of Python 3.7 compatibility, the minimum Pandas version will be updated to 1.5.0.
52-
if packaging.version.Version(pandas.__version__) >= packaging.version.Version("1.5.0"):
53-
from db_dtypes.json import JSONArray, JSONArrowType, JSONDtype
54-
else:
55-
JSONArray = None
56-
JSONDtype = None
57-
58-
5949
@pandas.api.extensions.register_extension_dtype
6050
class TimeDtype(core.BaseDatetimeDtype):
6151
"""
@@ -347,6 +337,22 @@ def __sub__(self, other):
347337
return super().__sub__(other)
348338

349339

340+
def _determine_all(json_array_type, json_dtype_type):
341+
"""Determines the list for __all__ based on JSON type availability."""
342+
base_all = [
343+
"__version__",
344+
"DateArray",
345+
"DateDtype",
346+
"TimeArray",
347+
"TimeDtype",
348+
]
349+
# Check if both JSON types are available (truthy)
350+
if json_array_type and json_dtype_type:
351+
return base_all + ["JSONDtype", "JSONArray", "JSONArrowType"]
352+
else:
353+
return base_all
354+
355+
350356
def _check_python_version():
351357
"""Checks the runtime Python version and issues a warning if needed."""
352358
sys_major, sys_minor, sys_micro = _versions_helpers.extract_runtime_version()
@@ -364,23 +370,4 @@ def _check_python_version():
364370

365371
_check_python_version()
366372

367-
368-
if not JSONArray or not JSONDtype:
369-
__all__ = [
370-
"__version__",
371-
"DateArray",
372-
"DateDtype",
373-
"TimeArray",
374-
"TimeDtype",
375-
]
376-
else:
377-
__all__ = [
378-
"__version__",
379-
"DateArray",
380-
"DateDtype",
381-
"JSONDtype",
382-
"JSONArray",
383-
"JSONArrowType",
384-
"TimeArray",
385-
"TimeDtype",
386-
]
373+
__all__ = _determine_all(JSONArray, JSONDtype)

tests/unit/test__init__.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,3 +83,62 @@ def test_check_python_version_does_not_warn_on_supported(mock_version_tuple):
8383

8484
# Assert that warnings.warn was NOT called
8585
mock_warn_call.assert_not_called()
86+
87+
88+
def test_determine_all_includes_json_when_available():
89+
"""
90+
Test that _determine_all includes JSON types when both are truthy.
91+
"""
92+
93+
from db_dtypes import _determine_all
94+
95+
# Simulate available types (can be any truthy object)
96+
mock_json_array = object()
97+
mock_json_dtype = object()
98+
99+
result = _determine_all(mock_json_array, mock_json_dtype)
100+
101+
expected_all = [
102+
"__version__",
103+
"DateArray",
104+
"DateDtype",
105+
"TimeArray",
106+
"TimeDtype",
107+
"JSONDtype",
108+
"JSONArray",
109+
"JSONArrowType",
110+
]
111+
assert set(result) == set(expected_all)
112+
assert "JSONDtype" in result
113+
assert "JSONArray" in result
114+
assert "JSONArrowType" in result
115+
116+
117+
@pytest.mark.parametrize(
118+
"mock_array, mock_dtype",
119+
[
120+
(None, object()), # JSONArray is None
121+
(object(), None), # JSONDtype is None
122+
(None, None), # Both are None
123+
],
124+
)
125+
def test_determine_all_excludes_json_when_unavailable(mock_array, mock_dtype):
126+
"""
127+
Test that _determine_all excludes JSON types if either is falsy.
128+
"""
129+
130+
from db_dtypes import _determine_all
131+
132+
result = _determine_all(mock_array, mock_dtype)
133+
134+
expected_all = [
135+
"__version__",
136+
"DateArray",
137+
"DateDtype",
138+
"TimeArray",
139+
"TimeDtype",
140+
]
141+
assert set(result) == set(expected_all)
142+
assert "JSONDtype" not in result
143+
assert "JSONArray" not in result
144+
assert "JSONArrowType" not in result

tests/unit/test_json.py

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import pytest
2222

2323
import db_dtypes
24+
import db_dtypes.json
2425

2526
# Check for minimum Pandas version.
2627
pytest.importorskip("pandas", minversion="1.5.0")
@@ -245,27 +246,45 @@ def cleanup_json_module_for_reload():
245246
# Explicitly register just in case it wasn't, or was cleaned up elsewhere.
246247
# This might raise ArrowKeyError itself if already registered, which is fine here.
247248
pa.register_extension_type(db_dtypes.json.JSONArrowType())
249+
248250
except pa.ArrowKeyError:
249251
pass # Already registered is the state we want before the test runs
250252

251253
# Remove the module from sys.modules so importlib.reload re-executes it
252-
if json_module_name in sys.modules: # COVERAGE FAIL: 252->255
254+
if json_module_name in sys.modules:
253255
del sys.modules[json_module_name]
254256

255257
yield # Run the test that uses this fixture
256258

257259
# Cleanup: Put the original module back if it existed
258260
# This helps isolate from other tests that might import db_dtypes.json
259-
if original_module: # COVERAGE FAIL: 259-261
261+
if original_module:
260262
sys.modules[json_module_name] = original_module
261-
elif json_module_name in sys.modules: # COVERAGE FAIL: 261->exit
263+
elif json_module_name in sys.modules:
262264
# If the test re-imported it but it wasn't there originally, remove it
263265
del sys.modules[json_module_name]
264266

265267
# Note: PyArrow doesn't have a public API to unregister types easily,
266268
# thus we are using the testing pattern of module isolation/reloading.
267269

268270

271+
# Test specifically for the fixture's pre-yield removal logic
272+
def test_fixture_removes_module_if_present(cleanup_json_module_for_reload):
273+
"""
274+
Tests that the cleanup_json_module_for_reload fixture removes
275+
db_dtypes.json from sys.modules before yielding to the test.
276+
This specifically targets the 'if json_module_name in sys.modules:' block.
277+
"""
278+
# This test runs *after* the fixture's `yield`.
279+
# The fixture should have removed the module if it was present.
280+
281+
json_module_name = "db_dtypes.json"
282+
283+
assert (
284+
json_module_name not in sys.modules
285+
), f"The fixture cleanup_json_module_for_reload should have removed {json_module_name}"
286+
287+
269288
def test_json_arrow_type_reregistration_is_handled(cleanup_json_module_for_reload):
270289
"""
271290
Verify that attempting to re-register JSONArrowType via module reload
@@ -276,17 +295,6 @@ def test_json_arrow_type_reregistration_is_handled(cleanup_json_module_for_reloa
276295
# forces Python to execute the module's top-level code again.
277296
# This includes the pa.register_extension_type call.
278297

279-
assert "db_dtypes.json" not in sys.modules
280-
try:
281-
import db_dtypes.json # noqa: F401
282-
283-
assert (
284-
True
285-
), "Module re-import completed without error, except block likely worked."
286-
287-
except pa.ArrowKeyError: # COVERAGE FAIL: 287-294
288-
# If this exception escapes, the except block in db_dtypes/json.py failed.
289-
pytest.fail(
290-
"pa.ArrowKeyError was raised during module reload, "
291-
"indicating the except block failed."
292-
)
298+
import db_dtypes.json # noqa: F401
299+
300+
assert True, "Module re-import completed without error, except block likely worked."

0 commit comments

Comments
 (0)