Skip to content

Commit 96a5e97

Browse files
committed
update json test for unittest coverage
1 parent 1bbd5f7 commit 96a5e97

File tree

1 file changed

+66
-0
lines changed

1 file changed

+66
-0
lines changed

tests/unit/test_json.py

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

15+
import importlib
1516
import json
17+
import sys
1618

1719
import numpy as np
1820
import pandas as pd
@@ -224,3 +226,67 @@ def test_json_arrow_record_batch():
224226
== '{"null_field":null,"order":{"address":{"city":"Anytown","street":"123 Main St"},"items":["book","pen","computer"],"total":15}}'
225227
)
226228
assert s[6] == "null"
229+
230+
@pytest.fixture
231+
def cleanup_json_module_for_reload():
232+
"""
233+
Fixture to ensure db_dtypes.json is registered and then removed
234+
from sys.modules to allow testing the registration except block via reload.
235+
"""
236+
json_module_name = "db_dtypes.json"
237+
original_module = sys.modules.get(json_module_name)
238+
239+
# 1. Ensure the type is registered initially (usually by the first import)
240+
try:
241+
# Make sure the module is loaded so the type exists
242+
import db_dtypes.json
243+
# Explicitly register just in case it wasn't, or was cleaned up elsewhere.
244+
# This might raise ArrowKeyError itself if already registered, which is fine here.
245+
pa.register_extension_type(db_dtypes.json.JSONArrowType())
246+
except pa.ArrowKeyError:
247+
pass # Already registered is the state we want before the test runs
248+
except ImportError:
249+
pytest.skip("Could not import db_dtypes.json to set up test.")
250+
251+
# 2. Remove the module from sys.modules so importlib.reload re-executes it
252+
if json_module_name in sys.modules:
253+
del sys.modules[json_module_name]
254+
255+
yield # Run the test that uses this fixture
256+
257+
# 3. Cleanup: Put the original module back if it existed
258+
# This helps isolate from other tests that might import db_dtypes.json
259+
if original_module:
260+
sys.modules[json_module_name] = original_module
261+
elif json_module_name in sys.modules:
262+
# If the test re-imported it but it wasn't there originally, remove it
263+
del sys.modules[json_module_name]
264+
265+
# Note: PyArrow doesn't have a public API to unregister types easily.
266+
# Relying on module isolation/reloading is a common testing pattern.
267+
268+
269+
def test_json_arrow_type_reregistration_is_handled(cleanup_json_module_for_reload):
270+
"""
271+
Verify that attempting to re-register JSONArrowType via module reload
272+
is caught by the except block and does not raise an error.
273+
"""
274+
try:
275+
# Re-importing the module after the fixture removed it from sys.modules
276+
# forces Python to execute the module's top-level code again.
277+
# This includes the pa.register_extension_type call.
278+
import db_dtypes.json
279+
280+
# If the import completes without raising pa.ArrowKeyError,
281+
# it means the 'except ArrowKeyError: pass' block worked as expected.
282+
assert True, "Module re-import completed without error, except block likely worked."
283+
284+
except pa.ArrowKeyError:
285+
# If this exception escapes, the except block in db_dtypes/json.py failed.
286+
pytest.fail(
287+
"pa.ArrowKeyError was raised during module reload, "
288+
"indicating the except block failed."
289+
)
290+
except Exception as e:
291+
# Catch any other unexpected error during the reload for better debugging.
292+
pytest.fail(f"An unexpected exception occurred during module reload: {e}")

0 commit comments

Comments
 (0)