Updates FutureWarning tests to account for unittest coverage

chalmerlowe · chalmerlowe · commit a0543b96e2fd · 2025-04-23T13:54:00.000Z
diff --git a/db_dtypes/__init__.py b/db_dtypes/__init__.py
@@ -339,7 +339,6 @@ def __sub__(self, other):
 
 
 sys_major, sys_minor, sys_micro = _versions_helpers.extract_runtime_version()
-print(f"DINOSAUR: {sys_major}.{sys_minor}.{sys_micro}")
 if sys_major == 3 and sys_minor in (7, 8):
     warnings.warn(
         "The python-bigquery library as well as the python-db-dtypes-pandas library no "
diff --git a/tests/unit/test__init__.py b/tests/unit/test__init__.py
@@ -1,5 +1,6 @@
 import sys
 import pytest
+import types
 import warnings
 from unittest import mock
 
@@ -9,198 +10,81 @@
 
 @pytest.fixture(autouse=True)
 def cleanup_imports():
+    """Ensures the target module and its helper are removed from sys.modules
+    before each test, allowing for clean imports with patching.
     """
-    Ensures the target module and its helper are removed from sys.modules
-    before and after each test, allowing for clean imports with patching.
-    """
-    # Store original sys.version_info if it's not already stored
-    if not hasattr(cleanup_imports, 'original_version_info'):
-        cleanup_imports.original_version_info = sys.version_info
-
-    # Remove modules before test
-    if MODULE_PATH in sys.modules:
-        del sys.modules[MODULE_PATH]
-    if HELPER_MODULE_PATH in sys.modules:
-         del sys.modules[HELPER_MODULE_PATH]
-
-    yield # Run the test
-
-    # Restore original sys.version_info after test
-    sys.version_info = cleanup_imports.original_version_info
-
-    # Remove modules after test
-    if MODULE_PATH in sys.modules:
-        del sys.modules[MODULE_PATH]
-    if HELPER_MODULE_PATH in sys.modules:
-         del sys.modules[HELPER_MODULE_PATH]
-
-
-@pytest.mark.parametrize(
-    "mock_version_tuple, version_str",
-    [
-        ((3, 7, 10), "3.7.10"),
-        ((3, 7, 0), "3.7.0"),
-        ((3, 8, 5), "3.8.5"),
-        ((3, 8, 12), "3.8.12"),
-    ]
-)
-def test_python_3_7_or_3_8_warning_on_import(mock_version_tuple, version_str):
-    """Test that a FutureWarning is raised for Python 3.7 during import."""
-    # Create a mock object mimicking sys.version_info attributes
-    # Use spec=sys.version_info to ensure it has the right attributes if needed,
-    # though just setting major/minor/micro is usually sufficient here.
-    mock_version_info = mock.Mock(spec=sys.version_info,
-                                  major=mock_version_tuple[0],
-                                  minor=mock_version_tuple[1],
-                                  micro=mock_version_tuple[2])
-
-    # Patch sys.version_info *before* importing db_dtypes
-    with mock.patch('sys.version_info', mock_version_info):
-        # Use pytest.warns to catch the expected warning during import
-        with pytest.warns(FutureWarning) as record:
-            # This import triggers __init__.py, which calls
-            # _versions_helpers.extract_runtime_version, which reads
-            # the *mocked* sys.version_info
-            import db_dtypes
-
-        # Assert that exactly one warning was recorded
-        assert len(record) == 1
-        warning_message = str(record[0].message)
-        # Assert the warning message content is correct
-        assert "longer supports Python 3.7 and Python 3.8" in warning_message
-
-@pytest.mark.parametrize(
-    "mock_version_tuple",
-    [
-        (3, 9, 1),    # Supported
-        (3, 10, 0),   # Supported
-        (3, 11, 2),   # Supported
-        (3, 12, 0),   # Supported
-    ]
-)
-def test_no_warning_for_other_versions_on_import(mock_version_tuple):
-    """Test that no FutureWarning is raised for other Python versions during import."""
-    with mock.patch(f"{MODULE_PATH}._versions_helpers.extract_runtime_version", return_value=mock_version_tuple):
-        # Use warnings.catch_warnings to check that NO relevant warning is raised
-        with warnings.catch_warnings(record=True) as record:
-            warnings.simplefilter("always") # Ensure warnings aren't filtered out by default config
-            import db_dtypes # Import triggers the code
-
-        # Assert that no FutureWarning matching the specific message was recorded
-        found_warning = False
-        for w in record:
-            # Check for the specific warning we want to ensure is NOT present
-            if (issubclass(w.category, FutureWarning) and
-                    "longer supports Python 3.7 and Python 3.8" in str(w.message)):
-                found_warning = True
-                break
-        assert not found_warning, f"Unexpected FutureWarning raised for Python version {mock_version_tuple}"
 
-
-@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.
-    """
-    # 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",
-    ]
+    # Store original modules that might exist
     original_modules = {}
-
-    # Store original modules and remove them
+    modules_to_clear = [MODULE_PATH, HELPER_MODULE_PATH]
     for mod_name in modules_to_clear:
-        original_modules[mod_name] = sys.modules.get(mod_name)
         if mod_name in sys.modules:
+            original_modules[mod_name] = sys.modules[mod_name]
             del sys.modules[mod_name]
 
     yield # Run the test
 
-    # Restore original modules after test
+    # Clean up again and restore originals if they existed
+    for mod_name in modules_to_clear:
+        if mod_name in sys.modules:
+            del sys.modules[mod_name] # Remove if test imported it
+    # Restore original modules
     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]
-
-
-# --- Test Case 1: JSON types available ---
-
-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__
-
-
-# --- Test Case 2: JSON types unavailable ---
 
 @pytest.mark.parametrize(
-    "patch_target_name",
+    "mock_version_tuple, version_str, expect_warning",
     [
-        "JSONArray",
-        "JSONDtype",
-        # Add both if needed, though one is sufficient to trigger the 'if'
-        # ("JSONArray", "JSONDtype"),
+        # Cases expected to warn
+        ((3, 7, 10), "3.7.10", True),
+        ((3, 7, 0), "3.7.0", True),
+        ((3, 8, 5), "3.8.5", True),
+        ((3, 8, 12), "3.8.12", True),
+        # Cases NOT expected to warn
+        ((3, 9, 1), "3.9.1", False),
+        ((3, 10, 0), "3.10.0", False),
+        ((3, 11, 2), "3.11.2", False),
+        ((3, 12, 0), "3.12.0", False),
     ]
 )
-def test_all_excludes_json_when_unavailable(cleanup_imports_for_all, patch_target_name):
-    """
-    Test that __all__ excludes JSON types when JSONArray or JSONDtype is unavailable (falsy).
+def test_python_version_warning_on_import(mock_version_tuple, version_str, expect_warning):
+    """Test that a FutureWarning is raised ONLY for Python 3.7 or 3.8 during import.
     """
-    patch_path = f"{MODULE_PATH}.{patch_target_name}"
-
-    # Patch one of the JSON types to be None *before* importing db_dtypes.
-    # This simulates the condition `if not JSONArray or not JSONDtype:` being true.
-    with mock.patch(patch_path, None):
-        # Need to ensure the json submodule itself is loaded if patching its contents
-        # If the patch target is directly in __init__, this isn't needed.
-        # Assuming JSONArray/JSONDtype are imported *into* __init__ from .json:
-        try:
-            import db_dtypes.json
-        except ImportError:
-             # Handle cases where the json module might genuinely be missing
-             pass
-
-        # Now import the main module, which will evaluate __all__
-        import db_dtypes
-
-        expected_all = [
-            "__version__",
-            "DateArray",
-            "DateDtype",
-            "TimeArray",
-            "TimeDtype",
-        ]
-        # Use set comparison for order independence
-        assert set(db_dtypes.__all__) == set(expected_all)
-        # Explicitly check absence of JSON types
-        assert "JSONDtype" not in db_dtypes.__all__
-        assert "JSONArray" not in db_dtypes.__all__
-        assert "JSONArrowType" not in db_dtypes.__all__
+    
+    # Create a mock function that returns the desired version tuple
+    mock_extract_func = mock.Mock(return_value=mock_version_tuple)
+
+    # Create a mock module object for _versions_helpers
+    mock_helpers_module = types.ModuleType(HELPER_MODULE_PATH)
+    mock_helpers_module.extract_runtime_version = mock_extract_func
+
+    # Use mock.patch.dict to temporarily replace the module in sys.modules
+    # This ensures that when db_dtypes.__init__ does `from . import _versions_helpers`,
+    # it gets our mock module.
+    with mock.patch.dict(sys.modules, {HELPER_MODULE_PATH: mock_helpers_module}):
+        if expect_warning:
+            with pytest.warns(FutureWarning) as record:
+                # The import will now use the mocked _versions_helpers module
+                import db_dtypes
+
+            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")
+                # The import will now use the mocked _versions_helpers module
+                import db_dtypes
+
+            found_warning = False
+            for w in record:
+                if (issubclass(w.category, FutureWarning) and
+                        "longer supports Python 3.7 and Python 3.8" in str(w.message)):
+                    found_warning = True
+                    break
+            assert not found_warning, (
+                f"Unexpected FutureWarning raised for Python version {version_str}"
+            )
