From 88064c07aaed654c6dde1adb8eae575e8f9d744a Mon Sep 17 00:00:00 2001
From: Borodin Gregory <grihabor@gmail.com>
Date: Wed, 22 Jan 2025 20:19:47 +0100
Subject: [PATCH] Add tests

---
 pyiceberg/table/update/__init__.py |  9 ++++++---
 tests/table/test_init.py           | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py
index 0914a2cfd0..0dfe3c9a0c 100644
--- a/pyiceberg/table/update/__init__.py
+++ b/pyiceberg/table/update/__init__.py
@@ -457,10 +457,13 @@ def _(update: SetSnapshotRefUpdate, base_metadata: TableMetadata, context: _Tabl
 
 @_apply_table_update.register(RemoveSnapshotsUpdate)
 def _(update: RemoveSnapshotsUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
-    if not any(s.snapshot_id == update.snapshot_id for s in base_metadata.snapshots):
-        raise ValueError(f"Snapshot with snapshot id {update.snapshot_id} does not exist")
+    for remove_snapshot_id in update.snapshot_ids:
+        if remove_snapshot_id == base_metadata.current_snapshot_id:
+            raise ValueError(f"Can't remove current snapshot id {remove_snapshot_id}")
+        if not any(s.snapshot_id == remove_snapshot_id for s in base_metadata.snapshots):
+            raise ValueError(f"Snapshot with snapshot id {remove_snapshot_id} does not exist: {base_metadata.snapshots}")
 
-    snapshots = [s for s in base_metadata.snapshots if s.snapshot_id != update.snapshot_id]
+    snapshots = [s for s in base_metadata.snapshots if s.snapshot_id not in update.snapshot_ids]
     context.add_update(update)
     return base_metadata.model_copy(update={"snapshots": snapshots})
 
diff --git a/tests/table/test_init.py b/tests/table/test_init.py
index e1f2ccc876..5d3e39e9ee 100644
--- a/tests/table/test_init.py
+++ b/tests/table/test_init.py
@@ -78,6 +78,7 @@
     AssertRefSnapshotId,
     AssertTableUUID,
     RemovePropertiesUpdate,
+    RemoveSnapshotsUpdate,
     RemoveStatisticsUpdate,
     SetDefaultSortOrderUpdate,
     SetPropertiesUpdate,
@@ -793,6 +794,33 @@ def test_update_metadata_set_snapshot_ref(table_v2: Table) -> None:
     )
 
 
+def test_update_remove_snapshots(table_v2: Table) -> None:
+    update = RemoveSnapshotsUpdate(
+        snapshot_ids=[3051729675574597004],
+    )
+    new_metadata = update_table_metadata(table_v2.metadata, (update,))
+    assert len(new_metadata.snapshots) == 1
+    assert new_metadata.snapshots[0].snapshot_id == 3055729675574597004
+    assert new_metadata.current_snapshot_id == 3055729675574597004
+    assert new_metadata.last_updated_ms > table_v2.metadata.last_updated_ms
+
+
+def test_update_remove_snapshots_doesnt_exist(table_v2: Table) -> None:
+    update = RemoveSnapshotsUpdate(
+        snapshot_ids=[123],
+    )
+    with pytest.raises(ValueError, match="Snapshot with snapshot id 123 does not exist"):
+        update_table_metadata(table_v2.metadata, (update,))
+
+
+def test_update_remove_snapshots_cant_remove_current_snapshot_id(table_v2: Table) -> None:
+    update = RemoveSnapshotsUpdate(
+        snapshot_ids=[3055729675574597004],
+    )
+    with pytest.raises(ValueError, match="Can't remove current snapshot id 3055729675574597004"):
+        update_table_metadata(table_v2.metadata, (update,))
+
+
 def test_update_metadata_add_update_sort_order(table_v2: Table) -> None:
     new_sort_order = SortOrder(order_id=table_v2.sort_order().order_id + 1)
     new_metadata = update_table_metadata(