add tests

Author: chinmay-bhat

tests/integration/test_snapshot_operations.py

@@ -40,3 +40,74 @@ def test_create_branch(catalog: Catalog) -> None:
     branch_snapshot_id = tbl.history()[-2].snapshot_id
     tbl.manage_snapshots().create_branch(snapshot_id=branch_snapshot_id, branch_name="branch123").commit()
     assert tbl.metadata.refs["branch123"] == SnapshotRef(snapshot_id=branch_snapshot_id, snapshot_ref_type="branch")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_manage_snapshots_context_manager(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
+    rollback_snapshot_id = tbl.history()[-4].snapshot_id
+    with tbl.manage_snapshots() as ms:
+        ms.create_tag(snapshot_id=current_snapshot_id, tag_name="testing")
+        ms.rollback_to_snapshot(snapshot_id=rollback_snapshot_id)
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
+    assert tbl.metadata.refs["testing"].snapshot_id == current_snapshot_id
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=rollback_snapshot_id, snapshot_ref_type="branch")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_rollback_to_snapshot(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    rollback_snapshot_id = tbl.history()[-3].snapshot_id
+    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
+    tbl.manage_snapshots().rollback_to_snapshot(snapshot_id=rollback_snapshot_id).commit()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=rollback_snapshot_id, snapshot_ref_type="branch")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_rollback_to_timestamp(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
+    timestamp = tbl.history()[-2].timestamp_ms
+    expected_snapshot_id = tbl.history()[-3].snapshot_id
+    # not inclusive of rollback_timestamp
+    tbl.manage_snapshots().rollback_to_timestamp(timestamp=timestamp).commit()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id, snapshot_ref_type="branch")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_set_current_snapshot_with_snapshot_id(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
+    expected_snapshot_id = tbl.history()[-3].snapshot_id
+    tbl.manage_snapshots().set_current_snapshot(snapshot_id=expected_snapshot_id).commit()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id, snapshot_ref_type="branch")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_set_current_snapshot_with_ref_name(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
+    expected_snapshot_id = tbl.history()[-3].snapshot_id
+    tbl.manage_snapshots().create_tag(snapshot_id=expected_snapshot_id, tag_name="test-tag19").commit()
+    tbl.manage_snapshots().set_current_snapshot(ref_name="test-tag19").commit()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id, snapshot_ref_type="branch")