|
4 | 4 | from __future__ import annotations |
5 | 5 |
|
6 | 6 | import json |
| 7 | +import multiprocessing as mp |
7 | 8 | from datetime import datetime |
8 | | -from unittest.mock import patch |
| 9 | +from multiprocessing.queues import Queue |
| 10 | +from multiprocessing.synchronize import Barrier |
| 11 | +from pathlib import Path |
| 12 | +from unittest.mock import MagicMock, patch |
9 | 13 |
|
10 | 14 | import pytest |
11 | 15 | from pyarrow import ArrowNotImplementedError |
@@ -502,3 +506,92 @@ def test_clear_partial_results_is_noop_when_no_partial_folder(tmp_path): |
502 | 506 | storage = ArtifactStorage(artifact_path=tmp_path) |
503 | 507 | assert not storage.partial_results_path.exists() |
504 | 508 | storage.clear_partial_results() # must not raise |
| 509 | + |
| 510 | + |
| 511 | +@pytest.mark.parametrize("existing", [False, True], ids=["fresh", "occupied"]) |
| 512 | +@patch("data_designer.engine.storage.artifact_storage.datetime") |
| 513 | +def test_create_dataset_directory_is_exclusive(mock_datetime: MagicMock, tmp_path: Path, existing: bool) -> None: |
| 514 | + mock_datetime.now.return_value = datetime(2025, 1, 1, 12, 3, 4) |
| 515 | + if existing: |
| 516 | + (tmp_path / "dataset").mkdir() |
| 517 | + (tmp_path / "dataset" / "sentinel").write_text("original") |
| 518 | + storages = [ArtifactStorage(artifact_path=tmp_path) for _ in range(3)] |
| 519 | + |
| 520 | + for storage in storages: |
| 521 | + storage.create_dataset_directory(resume=ResumeMode.NEVER) |
| 522 | + assert storage.base_dataset_path.is_dir() |
| 523 | + assert storage.media_storage.base_path == storage.base_dataset_path |
| 524 | + assert storage.media_storage.images_dir == storage.base_dataset_path / "images" |
| 525 | + |
| 526 | + assert len({storage.base_dataset_path for storage in storages}) == 3 |
| 527 | + assert len(list(tmp_path.iterdir())) == 3 + existing |
| 528 | + if existing: |
| 529 | + assert (tmp_path / "dataset" / "sentinel").read_text() == "original" |
| 530 | + |
| 531 | + |
| 532 | +def test_create_dataset_directory_does_not_reuse_empty_directory(tmp_path: Path) -> None: |
| 533 | + storage = ArtifactStorage(artifact_path=tmp_path) |
| 534 | + (tmp_path / "dataset").mkdir() |
| 535 | + |
| 536 | + storage.create_dataset_directory(resume=ResumeMode.NEVER) |
| 537 | + |
| 538 | + assert storage.base_dataset_path != tmp_path / "dataset" |
| 539 | + assert storage.base_dataset_path.is_dir() |
| 540 | + assert list((tmp_path / "dataset").iterdir()) == [] |
| 541 | + |
| 542 | + |
| 543 | +@pytest.mark.parametrize("resume", [ResumeMode.ALWAYS, ResumeMode.IF_POSSIBLE]) |
| 544 | +def test_create_dataset_directory_reuses_resumed_directory(tmp_path: Path, resume: ResumeMode) -> None: |
| 545 | + dataset_path = tmp_path / "dataset" |
| 546 | + dataset_path.mkdir() |
| 547 | + (dataset_path / "sentinel").write_text("original") |
| 548 | + storage = ArtifactStorage(artifact_path=tmp_path, resume=resume) |
| 549 | + |
| 550 | + storage.create_dataset_directory(resume=resume) |
| 551 | + |
| 552 | + assert storage.base_dataset_path == dataset_path |
| 553 | + assert list(tmp_path.iterdir()) == [dataset_path] |
| 554 | + assert (dataset_path / "sentinel").read_text() == "original" |
| 555 | + |
| 556 | + |
| 557 | +def test_create_dataset_directory_reports_filesystem_error(tmp_path: Path) -> None: |
| 558 | + storage = ArtifactStorage(artifact_path=tmp_path) |
| 559 | + with patch.object(Path, "mkdir", side_effect=PermissionError("permission denied")): |
| 560 | + with pytest.raises(ArtifactStorageError, match="Cannot create dataset directory"): |
| 561 | + storage.create_dataset_directory(resume=ResumeMode.NEVER) |
| 562 | + |
| 563 | + |
| 564 | +def _claim_dataset_directory(artifact_path: Path, barrier: Barrier, output: Queue) -> None: |
| 565 | + with patch("data_designer.engine.storage.artifact_storage.datetime") as clock: |
| 566 | + clock.now.return_value = datetime(2025, 1, 1, 12, 3, 4) |
| 567 | + storage = ArtifactStorage(artifact_path=artifact_path) |
| 568 | + barrier.wait(timeout=30) |
| 569 | + storage.create_dataset_directory(resume=ResumeMode.NEVER) |
| 570 | + output.put(str(storage.base_dataset_path)) |
| 571 | + |
| 572 | + |
| 573 | +@pytest.mark.parametrize("existing", [False, True], ids=["fresh", "occupied"]) |
| 574 | +def test_create_dataset_directory_is_exclusive_across_processes(tmp_path: Path, existing: bool) -> None: |
| 575 | + if existing: |
| 576 | + (tmp_path / "dataset").mkdir() |
| 577 | + (tmp_path / "dataset" / "sentinel").write_text("original") |
| 578 | + context = mp.get_context("spawn") |
| 579 | + barrier = context.Barrier(2) |
| 580 | + output = context.Queue() |
| 581 | + processes = [context.Process(target=_claim_dataset_directory, args=(tmp_path, barrier, output)) for _ in range(2)] |
| 582 | + try: |
| 583 | + for process in processes: |
| 584 | + process.start() |
| 585 | + for process in processes: |
| 586 | + process.join(timeout=45) |
| 587 | + assert process.exitcode == 0 |
| 588 | + paths = [Path(output.get(timeout=5)) for _ in processes] |
| 589 | + assert len(set(paths)) == 2 |
| 590 | + assert all(path.is_dir() for path in paths) |
| 591 | + finally: |
| 592 | + for process in processes: |
| 593 | + if process.is_alive(): |
| 594 | + process.terminate() |
| 595 | + process.join(timeout=5) |
| 596 | + output.close() |
| 597 | + output.join_thread() |
0 commit comments