import icechunk
import zarr
async def test_timetravel():
store = await icechunk.IcechunkStore.create(
storage=icechunk.StorageConfig.memory("test"),
config=icechunk.StoreConfig(inline_chunk_threshold_bytes=1),
)
group = zarr.group(store=store, overwrite=True)
air_temp = group.create_array(
"air_temp", shape=(1000, 1000), chunk_shape=(100, 100), dtype="i4"
)
air_temp[:, :] = 42
assert air_temp[200, 6] == 42
snapshot_id = await store.commit("commit 1")
air_temp[:, :] = 54
assert air_temp[200, 6] == 54
new_snapshot_id = await store.commit("commit 2")
await store.checkout(snapshot_id=snapshot_id)
assert air_temp[200, 6] == 42
await store.checkout(snapshot_id=new_snapshot_id)
assert air_temp[200, 6] == 54
await store.checkout(branch="main")
air_temp[:, :] = 76
assert store.has_uncommitted_changes
assert store.branch == "main"
assert store.snapshot_id == new_snapshot_id
await store.reset()
assert not store.has_uncommitted_changes
assert air_temp[200, 6] == 54
await store.new_branch("feature")
assert store.branch == "feature"
air_temp[:, :] = 90
feature_snapshot_id = await store.commit("commit 3")
await store.tag("v1.0", feature_snapshot_id)
await store.checkout(tag="v1.0")
assert store.branch is None
assert air_temp[200, 6] == 90
parents = [p async for p in store.ancestry()]
assert [snap.message for snap in parents] == [
"commit 3",
"commit 2",
"commit 1",
"Repository initialized",
]
assert sorted(parents, key=lambda p: p.written_at) == list(reversed(parents))
assert len(set([snap.id for snap in parents])) == 4