Skip to content

Commit 0e3c2b9

Browse files
committed
test: fold metadata-count test into dropped-rows test
Remove test_write_metadata_records_actual_and_target_counts (poked _actual_num_records directly) and assert metadata counts in test_dropped_rows_reduce_actual_record_count instead, which exercises the same path through the public API.
1 parent 7846f84 commit 0e3c2b9

1 file changed

Lines changed: 7 additions & 18 deletions

File tree

packages/data-designer-engine/tests/engine/dataset_builders/test_async_builder_integration.py

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -272,26 +272,10 @@ async def test_checkpoint_produces_correct_parquet_calls() -> None:
272272
# -- Partial completion warning ------------------------------------------------
273273

274274

275-
def test_write_metadata_records_actual_and_target_counts() -> None:
276-
"""write_metadata records the correct actual vs target counts."""
277-
storage = MagicMock()
278-
storage.dataset_name = "test"
279-
storage.get_file_paths.return_value = {}
280-
281-
buffer_manager = RowGroupBufferManager(storage)
282-
# Simulate 3 records checkpointed out of a 10-record target
283-
buffer_manager._actual_num_records = 3
284-
285-
buffer_manager.write_metadata(target_num_records=10, buffer_size=5)
286-
287-
written = storage.write_metadata.call_args[0][0]
288-
assert written["actual_num_records"] == 3
289-
assert written["target_num_records"] == 10
290-
291-
292275
@pytest.mark.asyncio(loop_scope="session")
293276
async def test_dropped_rows_reduce_actual_record_count() -> None:
294-
"""When all rows in a row group are dropped, actual_num_records reflects the shortfall."""
277+
"""When all rows in a row group are dropped, actual_num_records reflects the shortfall
278+
and write_metadata records the correct actual vs target counts."""
295279
provider = _mock_provider()
296280
seed_gen = MockSeed(config=_expr_config("seed"), resource_provider=provider)
297281

@@ -330,3 +314,8 @@ def drop_all_in_rg1(rg_id: int, rg_size: int) -> None:
330314
await scheduler.run()
331315

332316
assert buffer_manager.actual_num_records < num_records
317+
318+
buffer_manager.write_metadata(target_num_records=num_records, buffer_size=3)
319+
written = storage.write_metadata.call_args[0][0]
320+
assert written["actual_num_records"] == buffer_manager.actual_num_records
321+
assert written["target_num_records"] == num_records

0 commit comments

Comments
 (0)