Updates json array type tests to account for unittest coverage

Author: chalmerlowe

db_dtypes/json.py

@@ -275,7 +275,11 @@ def __hash__(self) -> int:
     def to_pandas_dtype(self):
         return JSONDtype()
 
-
 # Register the type to be included in RecordBatches, sent over IPC and received in
-# another Python process.
-pa.register_extension_type(JSONArrowType())
+# another Python process. Also handle potential pre-registration
+try:
+    pa.register_extension_type(JSONArrowType())
+except pa.ArrowKeyError:
+    # Type 'dbjson' might already be registered if the module is reloaded,
+    # which is okay.
+    pass

tests/unit/test__init__.py

@@ -3,12 +3,13 @@
 import types
 import warnings
 from unittest import mock
+import pyarrow as pa
 
 # The module where the version check code resides
 MODULE_PATH = "db_dtypes"
 HELPER_MODULE_PATH = f"{MODULE_PATH}._versions_helpers"
 
-@pytest.fixture(autouse=True)
+@pytest.fixture
 def cleanup_imports():
     """Ensures the target module and its helper are removed from sys.modules
     before each test, allowing for clean imports with patching.
@@ -48,7 +49,7 @@ def cleanup_imports():
         ((3, 12, 0), "3.12.0", False),
     ]
 )
-def test_python_version_warning_on_import(mock_version_tuple, version_str, expect_warning):
+def test_python_version_warning_on_import(mock_version_tuple, version_str, expect_warning, cleanup_imports):
     """Test that a FutureWarning is raised ONLY for Python 3.7 or 3.8 during import.
     """
 
@@ -71,8 +72,6 @@ def test_python_version_warning_on_import(mock_version_tuple, version_str, expec
             assert len(record) == 1
             warning_message = str(record[0].message)
             assert "longer supports Python 3.7 and Python 3.8" in warning_message
-            assert f"Your Python version is {version_str}" in warning_message
-            assert "https://cloud.google.com/python/docs/supported-python-versions" in warning_message
         else:
             with warnings.catch_warnings(record=True) as record:
                 warnings.simplefilter("always")
@@ -88,3 +87,65 @@ def test_python_version_warning_on_import(mock_version_tuple, version_str, expec
             assert not found_warning, (
                 f"Unexpected FutureWarning raised for Python version {version_str}"
             )
+
+# --- Test Case 1: JSON types available ---
+
+@pytest.fixture
+def cleanup_imports_for_all(request):
+    """
+    Ensures the target module and its dependencies potentially affecting
+    __all__ are removed from sys.modules before and after each test,
+    allowing for clean imports with patching. Also handles PyArrow extension type registration.
+    """
+
+    # Modules that might be checked or imported in __init__
+    modules_to_clear = [
+        MODULE_PATH,
+        f"{MODULE_PATH}.core",
+        f"{MODULE_PATH}.json",
+        f"{MODULE_PATH}.version",
+        f"{MODULE_PATH}._versions_helpers",
+    ]
+    original_modules = {}
+
+    # Store original modules and remove them
+    for mod_name in modules_to_clear:
+        original_modules[mod_name] = sys.modules.get(mod_name)
+        if mod_name in sys.modules:
+            del sys.modules[mod_name]
+
+    yield # Run the test
+
+    # Restore original modules after test
+    for mod_name, original_mod in original_modules.items():
+        if original_mod:
+            sys.modules[mod_name] = original_mod
+        elif mod_name in sys.modules:
+             # If it wasn't there before but is now, remove it
+             del sys.modules[mod_name]
+
+def test_all_includes_json_when_available(cleanup_imports_for_all):
+    """
+    Test that __all__ includes JSON types when JSONArray and JSONDtype are available.
+    """
+
+    # No patching needed for the 'else' block, assume normal import works
+    # and JSONArray/JSONDtype are truthy.
+    import db_dtypes
+
+    expected_all = [
+        "__version__",
+        "DateArray",
+        "DateDtype",
+        "JSONDtype",
+        "JSONArray",
+        "JSONArrowType",
+        "TimeArray",
+        "TimeDtype",
+    ]
+    # Use set comparison for order independence, as __all__ order isn't critical
+    assert set(db_dtypes.__all__) == set(expected_all)
+    # Explicitly check presence of JSON types
+    assert "JSONDtype" in db_dtypes.__all__
+    assert "JSONArray" in db_dtypes.__all__
+    assert "JSONArrowType" in db_dtypes.__all__