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,