Remove redundant concurrent snapshot expiration tests and improve thread safety assertions

Author: ForeverAngry

tests/table/test_expire_snapshots.py

@@ -15,9 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 import threading
-import uuid
 from datetime import datetime, timedelta
-from typing import Any, Dict, List
+from typing import Dict
 from unittest.mock import MagicMock, Mock
 from uuid import uuid4
 
@@ -283,324 +282,3 @@ def worker2() -> None:
 
     assert results["expire1_snapshots"] == expected_1, "Worker 1 snapshots contaminated"
     assert results["expire2_snapshots"] == expected_2, "Worker 2 snapshots contaminated"
-
-
-def test_concurrent_different_tables_expiration() -> None:
-    """Test that concurrent snapshot expiration on DIFFERENT tables works correctly.
-
-    This test reproduces the issue described in:
-    https://github.com/apache/iceberg-python/issues/2409
-
-    The issue occurs when expiring snapshots from different tables concurrently,
-    where snapshot IDs from one table get applied to another table.
-    """
-    # Create two mock tables with different snapshot IDs
-    table1 = Mock()
-    table1.metadata = Mock()
-    table1.metadata.table_uuid = uuid4()
-
-    table2 = Mock()
-    table2.metadata = Mock()
-    table2.metadata.table_uuid = uuid4()
-
-    # Track calls to each table's expire_snapshots method
-    table1_expire_calls = []
-    table2_expire_calls = []
-
-    def create_table1_expire_mock() -> Mock:
-        expire_mock = Mock()
-
-        def side_effect(sid: int) -> Mock:
-            table1_expire_calls.append(sid)
-            return expire_mock
-
-        expire_mock.by_id = Mock(side_effect=side_effect)
-        expire_mock.commit = Mock(return_value=None)
-        return expire_mock
-
-    def create_table2_expire_mock() -> Mock:
-        expire_mock = Mock()
-
-        def side_effect(sid: int) -> Mock:
-            table2_expire_calls.append(sid)
-            return expire_mock
-
-        expire_mock.by_id = Mock(side_effect=side_effect)
-        expire_mock.commit = Mock(return_value=None)
-        return expire_mock
-
-    table1.maintenance = Mock()
-    table1.maintenance.expire_snapshots = Mock(side_effect=create_table1_expire_mock)
-
-    table2.maintenance = Mock()
-    table2.maintenance.expire_snapshots = Mock(side_effect=create_table2_expire_mock)
-
-    # Define different snapshot IDs for each table
-    table1_snapshot_ids = [1001, 1002, 1003, 1004, 1005]
-    table2_snapshot_ids = [2001, 2002, 2003, 2004, 2005]
-
-    def expire_table_snapshots(table_obj: Any, table_name: str, snapshots_to_expire: List[int], results: Dict[str, Any]) -> None:
-        """Expire specific snapshots from a table."""
-        try:
-            # Expire the snapshots one by one (as in the user's example)
-            for snapshot_id in snapshots_to_expire:
-                table_obj.maintenance.expire_snapshots().by_id(snapshot_id).commit()
-
-            results["success"] = True
-            results["expired_snapshots"] = snapshots_to_expire
-
-        except Exception as e:
-            results["success"] = False
-            results["error"] = str(e)
-
-    # Prepare snapshots to expire (first 2 from each table)
-    table1_to_expire = table1_snapshot_ids[:2]
-    table2_to_expire = table2_snapshot_ids[:2]
-
-    results1: Dict[str, Any] = {}
-    results2: Dict[str, Any] = {}
-
-    # Create threads to expire snapshots from different tables concurrently
-    thread1 = threading.Thread(target=expire_table_snapshots, args=(table1, "table1", table1_to_expire, results1))
-    thread2 = threading.Thread(target=expire_table_snapshots, args=(table2, "table2", table2_to_expire, results2))
-
-    # Start threads concurrently
-    thread1.start()
-    thread2.start()
-
-    # Wait for completion
-    thread1.join()
-    thread2.join()
-
-    # Check results - both should succeed if thread safety is correct
-    # Assert both operations succeeded
-    assert results1.get("success", False), f"Table1 expiration failed: {results1.get('error', 'Unknown error')}"
-    assert results2.get("success", False), f"Table2 expiration failed: {results2.get('error', 'Unknown error')}"
-
-    # CRITICAL: Verify that each table only received its own snapshot IDs
-    # This is the key test - if the bug exists, snapshot IDs will cross-contaminate
-    for sid in table1_expire_calls:
-        assert sid in table1_snapshot_ids, f"Table1 received unexpected snapshot ID {sid}"
-
-    for sid in table2_expire_calls:
-        assert sid in table2_snapshot_ids, f"Table2 received unexpected snapshot ID {sid}"
-
-    # Verify expected snapshots were expired
-    assert set(table1_expire_calls) == set(table1_to_expire), "Table1 didn't expire expected snapshots"
-    assert set(table2_expire_calls) == set(table2_to_expire), "Table2 didn't expire expected snapshots"
-
-
-def test_concurrent_same_table_different_snapshots(table_v2_with_extensive_snapshots: Table) -> None:
-    """Test that concurrent snapshot expiration operations on the same table work correctly."""
-    # Mock the catalog's commit_table method for both operations
-    table_v2_with_extensive_snapshots.catalog = MagicMock()
-    table_v2_with_extensive_snapshots.catalog.commit_table.return_value = CommitTableResponse(
-        metadata=table_v2_with_extensive_snapshots.metadata, metadata_location="test://new_location"
-    )
-
-    # Use existing snapshot IDs from fixture data, but filter out protected snapshots
-    all_snapshots = list(table_v2_with_extensive_snapshots.snapshots())
-    snapshot_ids = [snapshot.snapshot_id for snapshot in all_snapshots]
-
-    # Get protected snapshot IDs from refs
-    protected_snapshot_ids = {ref.snapshot_id for ref in table_v2_with_extensive_snapshots.metadata.refs.values()}
-
-    # Find unprotected snapshots that we can expire
-    unprotected_snapshot_ids = [sid for sid in snapshot_ids if sid not in protected_snapshot_ids]
-
-    # If we don't have enough unprotected snapshots, skip the test
-    if len(unprotected_snapshot_ids) < 2:
-        pytest.skip("Not enough unprotected snapshots available for testing")
-
-    # We'll expire the first two unprotected snapshots concurrently
-    to_expire1 = [unprotected_snapshot_ids[0]]
-    to_expire2 = [unprotected_snapshot_ids[1]]
-
-    def expire_snapshots_thread_func(table: Any, snapshot_ids_to_expire: List[int], results: Dict[str, Any]) -> None:
-        """Function to run in a thread that expires snapshots and captures results."""
-        try:
-            # Expire snapshots
-            expire_op = table.maintenance.expire_snapshots()
-            for snapshot_id in snapshot_ids_to_expire:
-                expire_op = expire_op.by_id(snapshot_id)
-            expire_op.commit()
-            results["success"] = True
-        except Exception as e:
-            results["success"] = False
-            results["error"] = str(e)
-
-    # Prepare result dictionaries to capture thread outcomes
-    results1: Dict[str, Any] = {}
-    results2: Dict[str, Any] = {}
-
-    # Create threads to expire snapshots concurrently
-    thread1 = threading.Thread(
-        target=expire_snapshots_thread_func, args=(table_v2_with_extensive_snapshots, to_expire1, results1)
-    )
-    thread2 = threading.Thread(
-        target=expire_snapshots_thread_func, args=(table_v2_with_extensive_snapshots, to_expire2, results2)
-    )
-
-    # Start and join threads
-    thread1.start()
-    thread2.start()
-    thread1.join()
-    thread2.join()
-
-    # Assert that both operations succeeded
-    assert results1.get("success", False), f"Thread 1 expiration failed: {results1.get('error', 'Unknown error')}"
-    assert results2.get("success", False), f"Thread 2 expiration failed: {results2.get('error', 'Unknown error')}"
-
-    # Verify that both commit_table calls were made
-    assert table_v2_with_extensive_snapshots.catalog.commit_table.call_count == 2
-
-
-def test_cross_table_snapshot_id_isolation() -> None:
-    """Test that verifies snapshot IDs don't get mixed up between different tables.
-
-    This test validates the fix for GitHub issue #2409 by ensuring that concurrent
-    operations on different table objects properly isolate their snapshot IDs.
-    """
-
-    # Create two mock table objects to simulate the user's scenario
-    # Mock table 1 with its own snapshot IDs
-    table1 = Mock()
-    table1.metadata = Mock()
-    table1.metadata.table_uuid = uuid.uuid4()
-    table1_snapshot_ids = [1001, 1002, 1003, 1004, 1005]
-
-    # Mock table 2 with different snapshot IDs
-    table2 = Mock()
-    table2.metadata = Mock()
-    table2.metadata.table_uuid = uuid.uuid4()
-    table2_snapshot_ids = [2001, 2002, 2003, 2004, 2005]
-
-    # Track which snapshot IDs each table's expire operation receives
-    table1_expire_calls = []
-    table2_expire_calls = []
-
-    def mock_table1_expire() -> Mock:
-        expire_mock = Mock()
-
-        def side_effect(sid: int) -> Mock:
-            table1_expire_calls.append(sid)
-            return expire_mock
-
-        expire_mock.by_id = Mock(side_effect=side_effect)
-        expire_mock.commit = Mock(return_value=None)
-        return expire_mock
-
-    def mock_table2_expire() -> Mock:
-        expire_mock = Mock()
-
-        def side_effect(sid: int) -> Mock:
-            table2_expire_calls.append(sid)
-            return expire_mock
-
-        expire_mock.by_id = Mock(side_effect=side_effect)
-        expire_mock.commit = Mock(return_value=None)
-        return expire_mock
-
-    table1.maintenance = Mock()
-    table1.maintenance.expire_snapshots = Mock(side_effect=mock_table1_expire)
-    table2.maintenance = Mock()
-    table2.maintenance.expire_snapshots = Mock(side_effect=mock_table2_expire)
-
-    def expire_from_table(table: Any, table_name: str, snapshot_ids: List[int], results: Dict[str, Any]) -> None:
-        """Expire snapshots from a specific table."""
-        try:
-            for snapshot_id in snapshot_ids:
-                table.maintenance.expire_snapshots().by_id(snapshot_id).commit()
-            results["success"] = True
-            results["expired_ids"] = snapshot_ids
-        except Exception as e:
-            results["success"] = False
-            results["error"] = str(e)
-
-    # Prepare snapshots to expire
-    table1_to_expire = table1_snapshot_ids[:2]  # [1001, 1002]
-    table2_to_expire = table2_snapshot_ids[:2]  # [2001, 2002]
-
-    results1: Dict[str, Any] = {}
-    results2: Dict[str, Any] = {}
-
-    # Run concurrent expiration operations
-    thread1 = threading.Thread(target=expire_from_table, args=(table1, "table1", table1_to_expire, results1))
-    thread2 = threading.Thread(target=expire_from_table, args=(table2, "table2", table2_to_expire, results2))
-
-    thread1.start()
-    thread2.start()
-    thread1.join()
-    thread2.join()
-
-    # CRITICAL ASSERTION: Each table should only receive its own snapshot IDs
-    # If this fails, it means the thread safety bug exists
-
-    # Table1 should only see table1 snapshot IDs
-    assert all(sid in table1_snapshot_ids for sid in table1_expire_calls), (
-        f"Table1 received unexpected snapshot IDs: {table1_expire_calls} (should only contain {table1_snapshot_ids})"
-    )
-
-    # Table2 should only see table2 snapshot IDs
-    assert all(sid in table2_snapshot_ids for sid in table2_expire_calls), (
-        f"Table2 received unexpected snapshot IDs: {table2_expire_calls} (should only contain {table2_snapshot_ids})"
-    )
-
-    # Verify no cross-contamination
-    table1_received_table2_ids = [sid for sid in table1_expire_calls if sid in table2_snapshot_ids]
-    table2_received_table1_ids = [sid for sid in table2_expire_calls if sid in table1_snapshot_ids]
-
-    assert len(table1_received_table2_ids) == 0, f"Table1 incorrectly received Table2 snapshot IDs: {table1_received_table2_ids}"
-
-    assert len(table2_received_table1_ids) == 0, f"Table2 incorrectly received Table1 snapshot IDs: {table2_received_table1_ids}"
-
-
-def test_batch_expire_snapshots(table_v2_with_extensive_snapshots: Table) -> None:
-    """Test that batch expiration of multiple snapshots works correctly."""
-    # Mock the catalog's commit_table method
-    table_v2_with_extensive_snapshots.catalog = MagicMock()
-    table_v2_with_extensive_snapshots.catalog.commit_table.return_value = CommitTableResponse(
-        metadata=table_v2_with_extensive_snapshots.metadata, metadata_location="test://new_location"
-    )
-
-    # Use existing snapshot IDs from fixture data, but filter out protected snapshots
-    all_snapshots = list(table_v2_with_extensive_snapshots.snapshots())
-    snapshot_ids = [snapshot.snapshot_id for snapshot in all_snapshots]
-
-    # Get protected snapshot IDs from refs
-    protected_snapshot_ids = {ref.snapshot_id for ref in table_v2_with_extensive_snapshots.metadata.refs.values()}
-
-    # Find unprotected snapshots that we can expire
-    unprotected_snapshot_ids = [sid for sid in snapshot_ids if sid not in protected_snapshot_ids]
-
-    # If we don't have enough unprotected snapshots, skip the test
-    if len(unprotected_snapshot_ids) < 2:
-        pytest.skip("Not enough unprotected snapshots available for testing")
-
-    # We'll expire the first two unprotected snapshots in a batch
-    to_expire = unprotected_snapshot_ids[:2]
-
-    def batch_expire_thread_func(table: Any, snapshot_ids_to_expire: List[int], results: Dict[str, Any]) -> None:
-        try:
-            # Expire all snapshots in a single batch operation
-            table.maintenance.expire_snapshots().by_ids(snapshot_ids_to_expire).commit()
-            results["success"] = True
-        except Exception as e:
-            results["success"] = False
-            results["error"] = str(e)
-
-    # Prepare result dictionary to capture thread outcome
-    results: Dict[str, Any] = {}
-
-    # Create thread to expire snapshots
-    thread = threading.Thread(target=batch_expire_thread_func, args=(table_v2_with_extensive_snapshots, to_expire, results))
-
-    # Start and join thread
-    thread.start()
-    thread.join()
-
-    # Assert that the operation succeeded
-    assert results.get("success", False), f"Batch expiration failed: {results.get('error', 'Unknown error')}"
-
-    # Verify that commit_table was called once
-    assert table_v2_with_extensive_snapshots.catalog.commit_table.call_count == 1