PROD-421: Reset generation callback state each time new generation st…

…arts.

GitOrigin-RevId: e1a6462647ef28b1fd3955c25b5314d6f293db31

src/gretel_synthetics/batch.py

@@ -865,6 +865,8 @@ def generate_all(
         progress_callback = None
         if callback:
             progress_callback = _GenerationCallback(callback, callback_interval)
+            # reset the underlying callback, since we are starting a new genaration
+            callback(GenerationProgress(), reset=True)
 
         self.reset()
         if output is not None and output not in ("df",):

tests/test_generate.py

@@ -1,6 +1,6 @@
 import threading
 
-from typing import Iterator, List
+from typing import Any, Iterator, List
 from unittest.mock import Mock
 
 import pytest
@@ -89,6 +89,32 @@ def test_generate_doesnt_return_partial_record_when_stopped(tmp_path):
     assert record is None
 
 
+def test_generate_resets_previous_progress(tmp_path):
+    dummy_dir = str(tmp_path)
+
+    rf = RecordFactory(
+        num_lines=10,
+        batches={},
+        header_list=["colA", "colB"],
+        delimiter="|",
+    )
+
+    reset_called = False
+
+    def callback(_: Any, *, reset=False):
+        if reset is True:
+            nonlocal reset_called
+            reset_called = True
+
+    result = rf.generate_all(
+        callback=callback, callback_interval=2, callback_threading=True
+    )
+
+    assert reset_called is True
+    # all will be empty, as there are no batches configured
+    assert result.records == [{}] * 10
+
+
 def _gen_and_set_thread_event(factory: RecordFactory) -> Iterator[gen_text]:
     """
     This generator: