3939from icechunk import (
4040 IcechunkError ,
4141 Repository ,
42+ RepositoryConfig ,
4243 SnapshotInfo ,
4344 Storage ,
4445 in_memory_storage ,
4546)
47+ from icechunk .testing .strategies import repository_configs
4648from zarr .testing .stateful import SyncStoreWrapper
4749
4850# JSON file contents, keep it simple
@@ -110,7 +112,10 @@ class TagModel:
110112
111113class Model :
112114 def __init__ (self , ** kwargs : Any ) -> None :
113- self .store : dict [str , Any ] = {} #
115+ self .store : dict [str , Any ] = {}
116+
117+ self .spec_version = 1 # will be overwritten on `@initialize`
118+ self .num_updates : int = 0
114119
115120 self .initial_snapshot_id : str | None = None
116121 self .changes_made : bool = False
@@ -151,6 +156,9 @@ def __setitem__(self, key: str, value: Buffer) -> None:
151156 def __getitem__ (self , key : str ) -> Buffer :
152157 return cast (Buffer , self .store [key ])
153158
159+ def upgrade (self ) -> None :
160+ self .num_updates += 1
161+
154162 @property
155163 def has_commits (self ) -> bool :
156164 return bool (self .commits )
@@ -171,10 +179,11 @@ def commit(self, snap: SnapshotInfo) -> None:
171179
172180 assert self .branch is not None
173181 self .branch_heads [self .branch ] = ref
182+ self .num_updates += 1
174183
175184 def amend (self , snap : SnapshotInfo ) -> None :
176185 """Amend the HEAD commit."""
177- # this is simpe because we aren't modeling the branch as a list of commits
186+ # this is simple because we aren't modeling the branch as a list of commits
178187 self .commit (snap )
179188
180189 def checkout_commit (self , ref : str ) -> None :
@@ -189,6 +198,7 @@ def checkout_commit(self, ref: str) -> None:
189198 def create_branch (self , name : str , commit : str ) -> None :
190199 assert commit in self .commits
191200 self .branch_heads [name ] = commit
201+ self .num_updates += 1
192202
193203 def checkout_branch (self , ref : str ) -> None :
194204 self .checkout_commit (self .branch_heads [ref ])
@@ -198,17 +208,27 @@ def checkout_branch(self, ref: str) -> None:
198208 def reset_branch (self , branch : str , commit : str ) -> None :
199209 assert commit in self .commits
200210 self .branch_heads [branch ] = commit
211+ self .num_updates += 1
201212
202213 def delete_branch (self , branch_name : str ) -> None :
214+ self ._delete_branch (branch_name )
215+ self .num_updates += 1
216+
217+ def _delete_branch (self , branch_name : str ) -> None :
203218 del self .branch_heads [branch_name ]
204219
205220 def delete_tag (self , tag : str ) -> None :
221+ self ._delete_tag (tag )
222+ self .num_updates += 1
223+
224+ def _delete_tag (self , tag : str ) -> None :
206225 del self .tags [tag ]
207226
208227 def create_tag (self , tag_name : str , commit_id : str ) -> None :
209228 assert commit_id in self .commits
210229 self .tags [tag_name ] = TagModel (commit_id = str (commit_id ))
211230 self .created_tags .add (tag_name )
231+ self .num_updates += 1
212232
213233 def checkout_tag (self , ref : str ) -> None :
214234 self .checkout_commit (self .tags [str (ref )].commit_id )
@@ -261,7 +281,7 @@ def expire_snapshots(
261281 }
262282 note (f"deleting tags { tags_to_delete = !r} " )
263283 for tag in tags_to_delete :
264- self .delete_tag (tag )
284+ self ._delete_tag (tag )
265285 else :
266286 tags_to_delete = set ()
267287
@@ -274,10 +294,12 @@ def expire_snapshots(
274294 note (f"deleting branches { branches_to_delete = !r} " )
275295 for branch in branches_to_delete :
276296 note (f"deleting { branch = !r} , { self .branch_heads [branch ]= !r} " )
277- self .delete_branch (branch )
297+ self ._delete_branch (branch )
278298 else :
279299 branches_to_delete = set ()
280300
301+ self .num_updates += 1
302+
281303 return ExpireInfo (
282304 expired_snapshots = expired_snaps ,
283305 deleted_branches = branches_to_delete ,
@@ -301,6 +323,7 @@ def garbage_collect(self, older_than: datetime.datetime) -> set[str]:
301323 self .ondisk_snaps .pop (k , None )
302324 deleted .add (k )
303325 note (f"Deleted snapshots in model: { deleted !r} " )
326+ self .num_updates += 1
304327 return deleted
305328
306329
@@ -325,7 +348,14 @@ def __init__(self) -> None:
325348 @initialize (data = st .data (), target = branches , spec_version = st .sampled_from ([1 , 2 ]))
326349 def initialize (self , data : st .DataObject , spec_version : Literal [1 , 2 ]) -> str :
327350 self .storage = in_memory_storage ()
328- self .repo = Repository .create (self .storage , spec_version = spec_version )
351+ config = data .draw (repository_configs ())
352+ self .model .spec_version = spec_version
353+
354+ self .repo = Repository .create (
355+ self .storage ,
356+ spec_version = spec_version ,
357+ config = config ,
358+ )
329359 self .session = self .repo .writable_session (DEFAULT_BRANCH )
330360
331361 snap = next (iter (self .repo .ancestry (branch = DEFAULT_BRANCH )))
@@ -338,6 +368,9 @@ def initialize(self, data: st.DataObject, spec_version: Literal[1, 2]) -> str:
338368 self .model .HEAD = HEAD
339369 self .model .create_branch (DEFAULT_BRANCH , HEAD )
340370 self .model .checkout_branch (DEFAULT_BRANCH )
371+ # RepoInitializedUpdate includes the initial branch creation,
372+ # so reset to 1 after create_branch incremented it.
373+ self .model .num_updates = 1
341374
342375 # initialize with some data always
343376 # TODO: always setting array metadata, since we cannot overwrite an existing group's zarr.json
@@ -371,17 +404,22 @@ def upgrade_spec_version(self) -> None:
371404 # don't test simple cases of catching error upgradging a v2 spec
372405 # that should be covered in unit tests
373406 icechunk .upgrade_icechunk_repository (self .repo )
407+ self .model .upgrade ()
374408 # TODO: remove the reopen after https://github.com/earth-mover/icechunk/issues/1521
375- self .reopen_repository ()
409+ self ._reopen_repository ()
376410
377- @rule ()
378- def reopen_repository (self ) -> None :
411+ @rule (data = st .data ())
412+ def reopen_repository (self , data : st .DataObject ) -> None :
413+ config = data .draw (repository_configs ())
414+ self ._reopen_repository (config )
415+
416+ def _reopen_repository (self , config : RepositoryConfig | None = None ) -> None :
379417 """Reopen the repository from storage to get fresh state.
380418
381419 This discards any uncommitted changes.
382420 """
383421 assert self .storage is not None , "storage must be initialized"
384- self .repo = Repository .open (self .storage )
422+ self .repo = Repository .open (self .storage , config = config )
385423 note (f"Reopened repository (spec_version={ self .repo .spec_version } )" )
386424
387425 # Reopening discards uncommitted changes - reset model to last committed state
@@ -688,6 +726,17 @@ def check_commit(self, commit: str) -> None:
688726 # even after expiration, written_at is unmodified
689727 assert actual .written_at == expected .written_at
690728
729+ def check_ops_log (self ) -> None :
730+ if self .model .spec_version == 1 :
731+ return
732+ actual_ops = list (self .repo .ops_log ())
733+ assert len (actual_ops ) == self .model .num_updates , (
734+ actual_ops ,
735+ self .model .num_updates ,
736+ actual_ops ,
737+ )
738+ assert isinstance (actual_ops [- 1 ], icechunk .RepoInitializedUpdate )
739+
691740 @invariant ()
692741 def checks (self ) -> None :
693742 # this method only exists to reduce verbosity of hypothesis output
@@ -697,6 +746,7 @@ def checks(self) -> None:
697746 self .check_tags ()
698747 self .check_branches ()
699748 self .check_ancestry ()
749+ self .check_ops_log ()
700750
701751 def check_list_prefix_from_root (self ) -> None :
702752 model_list = self .model .list_prefix ("" )
0 commit comments