"""Unit tests for the Iceberg lake sync: incremental tile diffing and add_files planning against a local filesystem fixture (a miniature lake in a tmp dir, a file:// warehouse, a sqlite catalog). No network.""" import datetime import os import sys import pyarrow as pa import pyarrow.parquet as pq import pytest from fsspec.implementations.local import LocalFileSystem sys.path.insert(0, os.path.dirname(__file__)) import sync_iceberg as si # noqa: E402 # Mirrors the real part-file schema (float32/float64 split, int16 doy, # all-null float64 sun) so the compatibility check is exercised for real. PART_SCHEMA = pa.schema([ ("date", pa.date32()), ("tmax", pa.float32()), ("tmin", pa.float32()), ("precip", pa.float32()), ("wind", pa.float32()), ("gust", pa.float32()), ("humid", pa.float32()), ("fmax", pa.float32()), ("fmin", pa.float64()), ("sun", pa.float64()), ("feels", pa.float64()), ("doy", pa.int16()), ("lat_idx", pa.int32()), ("lon_idx", pa.int32()), ]) POINTS = 2 # points per fixture tile DAYS = 3 # rows per point per month def write_tile(root, ti, tj, months): """Write one tile's hive part files; return its total row count.""" total = 0 for year, month in months: rows = {name: [] for name in PART_SCHEMA.names} for p in range(POINTS): i, j = ti * si.TILE + p, tj * si.TILE + p for d in range(DAYS): day = datetime.date(year, month, d + 1) rows["date"].append(day) for col in ("tmax", "tmin", "precip", "wind", "gust", "humid", "fmax", "fmin", "feels"): rows[col].append(float(10 * p + d)) rows["sun"].append(None) rows["doy"].append(day.timetuple().tm_yday) rows["lat_idx"].append(i) rows["lon_idx"].append(j) part = pa.table(rows, schema=PART_SCHEMA) path = os.path.join(root, si.PREFIX, "daily", f"tile={ti}_{tj}", f"year={year}", f"month={month:02d}", "part.parquet") os.makedirs(os.path.dirname(path), exist_ok=True) pq.write_table(part, path) total += part.num_rows return total def write_manifest(root, tile_rows): """Manifest with one row per (tile, point) like the extractor writes; the sync only reads tile+rows.""" tiles, rows = [], [] for tile, total in tile_rows.items(): per_point = total // POINTS for _ in range(POINTS): tiles.append(tile) rows.append(per_point) m = pa.table({ "lat_idx": pa.array([0] * len(tiles), pa.int64()), "lon_idx": pa.array([0] * len(tiles), pa.int64()), "lat": pa.array([0.0] * len(tiles)), "lon": pa.array([0.0] * len(tiles)), "tile": pa.array(tiles, pa.large_string()), "rows": pa.array(rows, pa.int64()), "date_min": pa.array(["1940-01-01"] * len(tiles), pa.large_string()), "date_max": pa.array(["2026-01-01"] * len(tiles), pa.large_string()), }) path = os.path.join(root, si.MANIFEST_KEY) os.makedirs(os.path.dirname(path), exist_ok=True) pq.write_table(m, path) @pytest.fixture def lake(tmp_path): root = str(tmp_path / "lake") counts = { "10_10": write_tile(root, 10, 10, [(2000, 1), (2000, 2)]), "10_11": write_tile(root, 10, 11, [(2000, 1)]), } write_manifest(root, counts) hl = si.HiveLake(LocalFileSystem(), root, f"file://{root}", None) hl.fixture_counts = counts hl.fixture_root = root return hl @pytest.fixture def catalog_db(tmp_path): return str(tmp_path / "catalog.db") def test_pending_tiles_diff(): manifest = {"10_10": 1, "10_11": 1, "2_3": 1} assert si.pending_tiles(manifest, set()) == ["2_3", "10_10", "10_11"] assert si.pending_tiles(manifest, {"10_10"}) == ["2_3", "10_11"] assert si.pending_tiles(manifest, set(manifest)) == [] # numeric, not lexicographic, tile ordering assert si.pending_tiles({"10_2": 1, "9_11": 1}, set()) == ["9_11", "10_2"] def test_initial_sync_registers_all_files_in_place(lake, catalog_db): result = si.run(lake, catalog_db, batch=1) expected = sum(lake.fixture_counts.values()) assert result["added_records"] == expected assert result["total_records"] == expected assert result["added_files"] == 3 # 2 months + 1 month assert not result["mismatches"] table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER) assert si.synced_tiles(table) == set(lake.fixture_counts) # add_files, not a rewrite: every data file is an original hive part file files = table.inspect.files().column("file_path").to_pylist() assert len(files) == 3 assert all(f"/{si.PREFIX}/daily/tile=" in f for f in files) # nothing was written outside iceberg/ (the hive tree is untouched, the # table tree holds only metadata) fs = LocalFileSystem() written = fs.find(os.path.join(lake.fixture_root, si.ICEBERG_PREFIX)) assert written and all("/metadata/" in p for p in written) def test_rerun_is_noop(lake, catalog_db): first = si.run(lake, catalog_db) again = si.run(lake, catalog_db) assert again["pending"] == [] assert again["added_records"] == 0 assert again["total_records"] == first["total_records"] assert again["metadata_location"] == first["metadata_location"] def test_incremental_new_tile_only(lake, catalog_db): si.run(lake, catalog_db) counts = dict(lake.fixture_counts) counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)]) write_manifest(lake.fixture_root, counts) result = si.run(lake, catalog_db) assert result["pending"] == ["11_10"] assert result["added_records"] == counts["11_10"] assert result["total_records"] == sum(counts.values()) def test_version_hint_tracks_current_metadata(lake, catalog_db): first = si.run(lake, catalog_db) hint_path = os.path.join(lake.fixture_root, si.ICEBERG_PREFIX, si.TABLE_NAME, "metadata", "version-hint.text") with open(hint_path) as f: assert f.read() == os.path.basename( first["metadata_location"]).removesuffix(".metadata.json") counts = dict(lake.fixture_counts) counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)]) write_manifest(lake.fixture_root, counts) second = si.run(lake, catalog_db) assert second["metadata_location"] != first["metadata_location"] with open(hint_path) as f: assert f.read() == os.path.basename( second["metadata_location"]).removesuffix(".metadata.json") def test_partition_spec_and_pruned_scan(lake, catalog_db): si.run(lake, catalog_db) table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER) assert [f.name for f in table.spec().fields] == \ ["tile_lat", "tile_lon", "year", "month"] parts = table.inspect.partitions() keys = {(r["partition"]["tile_lat"], r["partition"]["tile_lon"]) for r in parts.to_pylist()} assert keys == {(120, 120), (120, 132)} # ti*12, tj*12 # a tile-shaped predicate resolves to that tile's rows only scan = table.scan(row_filter="lat_idx >= 120 and lat_idx < 132 " "and lon_idx >= 132 and lon_idx < 144") assert scan.to_arrow().num_rows == lake.fixture_counts["10_11"] # and file planning pruned to that tile's single part file assert len(list(scan.plan_files())) == 1 def test_lost_catalog_db_reregisters_without_readding(lake, catalog_db): first = si.run(lake, catalog_db) os.remove(catalog_db) # simulate losing local state again = si.run(lake, catalog_db) # synced-tiles lives in table metadata, so nothing is re-added assert again["pending"] == [] assert again["added_records"] == 0 assert again["total_records"] == first["total_records"] def test_dry_run_writes_nothing(lake, catalog_db): result = si.run(lake, catalog_db, dry_run=True) assert sorted(result["pending"]) == sorted(lake.fixture_counts) assert result["metadata_location"] is None assert not os.path.exists( os.path.join(lake.fixture_root, si.ICEBERG_PREFIX)) def test_manifest_tile_without_files_refuses(lake, catalog_db): counts = dict(lake.fixture_counts) counts["50_50"] = 6 # in the manifest, no files write_manifest(lake.fixture_root, counts) # an invariant break aborts immediately: no retries, nothing marked synced with pytest.raises(si.LakeInvariantError, match="50_50"): si.run(lake, catalog_db) def test_transient_failure_retries_and_succeeds(lake, catalog_db, monkeypatch): monkeypatch.setattr(si, "BACKOFF_SECONDS", 0) fails = {"n": 2} original = si.HiveLake.tile_files def flaky(self, tile): if fails["n"]: fails["n"] -= 1 raise OSError("AWS Error SLOW_DOWN during CompleteMultipartUpload") return original(self, tile) monkeypatch.setattr(si.HiveLake, "tile_files", flaky) result = si.run(lake, catalog_db, retries=2) assert result["added_records"] == sum(lake.fixture_counts.values()) assert not result["mismatches"] def test_transient_failure_exhausts_retries(lake, catalog_db, monkeypatch): monkeypatch.setattr(si, "BACKOFF_SECONDS", 0) monkeypatch.setattr( si.HiveLake, "tile_files", lambda self, tile: (_ for _ in ()).throw(OSError("SLOW_DOWN"))) with pytest.raises(OSError): si.run(lake, catalog_db, retries=1)