12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
+ import importlib
15
16
import json
17
+ import sys
16
18
17
19
import numpy as np
18
20
import pandas as pd
@@ -224,3 +226,67 @@ def test_json_arrow_record_batch():
224
226
== '{"null_field":null,"order":{"address":{"city":"Anytown","street":"123 Main St"},"items":["book","pen","computer"],"total":15}}'
225
227
)
226
228
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