From 4f6faa91f6f7ba896f0ddb66170d6b6b962a443c Mon Sep 17 00:00:00 2001
From: Yingjian Wu <yingjianw@netflix.com>
Date: Mon, 28 Oct 2024 15:40:25 -0700
Subject: [PATCH] address comments

---
 tests/catalog/test_base.py                   | 23 --------------------
 tests/integration/test_writes/test_writes.py | 23 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/tests/catalog/test_base.py b/tests/catalog/test_base.py
index df64a7b2c6..e212854ee2 100644
--- a/tests/catalog/test_base.py
+++ b/tests/catalog/test_base.py
@@ -766,26 +766,3 @@ def test_table_properties_raise_for_none_value(catalog: InMemoryCatalog) -> None
     with pytest.raises(ValidationError) as exc_info:
         _ = given_catalog_has_a_table(catalog, properties=property_with_none)
     assert "None type is not a supported value in properties: property_name" in str(exc_info.value)
-
-
-def test_abort_table_transaction_on_exception(catalog: InMemoryCatalog) -> None:
-    tbl = given_catalog_has_a_table(catalog)
-    # Populate some initial data
-    data = pa.Table.from_pylist(
-        [{"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}],
-        schema=TEST_TABLE_SCHEMA.as_arrow(),
-    )
-    tbl.append(data)
-
-    # Data to overwrite
-    data = pa.Table.from_pylist(
-        [{"x": 7, "y": 8, "z": 9}, {"x": 7, "y": 8, "z": 9}, {"x": 7, "y": 8, "z": 9}],
-        schema=TEST_TABLE_SCHEMA.as_arrow(),
-    )
-
-    with pytest.raises(ValueError):
-        with tbl.transaction() as txn:
-            txn.overwrite(data)
-            raise ValueError
-
-    assert len(tbl.scan().to_pandas()) == 2  # type: ignore
diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py
index fc2746c614..c1024193d2 100644
--- a/tests/integration/test_writes/test_writes.py
+++ b/tests/integration/test_writes/test_writes.py
@@ -1448,3 +1448,26 @@ def test_rewrite_manifest_after_partition_evolution(session_catalog: Catalog) ->
             EqualTo("category", "A"),
         ),
     )
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_abort_table_transaction_on_exception(
+    spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int
+) -> None:
+    identifier = "default.table_test_abort_table_transaction_on_exception"
+    tbl = _create_table(session_catalog, identifier, properties={"format-version": format_version})
+
+    # Pre-populate some data
+    tbl.append(arrow_table_with_null)
+    assert len(tbl.scan().to_pandas()) == 3
+
+    # try to commit a transaction that raises exception at the middle
+    with pytest.raises(ValueError):
+        with tbl.transaction() as txn:
+            txn.append(arrow_table_with_null)
+            raise ValueError
+            txn.append(arrow_table_with_null)  # type: ignore
+
+    # Validate the transaction is aborted
+    assert len(tbl.scan().to_pandas()) == 3  # type: ignore