diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst
index ef75eef..ccc18bc 100644
--- a/doc/source/changelog.rst
+++ b/doc/source/changelog.rst
@@ -5,7 +5,7 @@ Changelog
 
 X.Y.Z (DD-MM-YYYY)
 ------------------
-* Add test case for irregular grids (:pr:`39`)
+* Add test case for irregular grids (:pr:`39`, :pr:`40`)
 * Rename MSv2PartitionEntryPoint to MSv2EntryPoint (:pr:`38`)
 * Move ``chunks`` kwarg functionality in MSv2PartitionEntryPoint.open_datatree
   to ``partition_chunks`` (:pr:`37`)
diff --git a/xarray_ms/testing/simulator.py b/xarray_ms/testing/simulator.py
index 8d4f821..b133777 100644
--- a/xarray_ms/testing/simulator.py
+++ b/xarray_ms/testing/simulator.py
@@ -112,6 +112,7 @@ class MSStructureSimulator:
   simulate_data: bool
   model: Dict[str, Any]
   data_description: DataDescription
+  transform_data: Callable[[PartitionDataType], PartitionDataType] | None
 
   def __init__(
     self,