46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
import asyncio
|
|
from contextlib import redirect_stderr, redirect_stdout
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
|
|
from assembly_dataset.config import Config
|
|
from assembly_dataset.pipeline import generate, export
|
|
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def test_generation_export_and_resume_report_exact_total(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=7)
|
|
output, errors = StringIO(), StringIO()
|
|
with redirect_stdout(output), redirect_stderr(errors):
|
|
asyncio.run(generate(config, show_progress=True))
|
|
report = export(config, show_progress=True)
|
|
asyncio.run(generate(config, show_progress=True))
|
|
self.assertEqual(report['examples'], 7)
|
|
self.assertEqual(output.getvalue(), '')
|
|
text = errors.getvalue()
|
|
self.assertIn('Генерация', text)
|
|
self.assertIn('Экспорт', text)
|
|
self.assertIn('7/7', text.replace(' ', ''))
|
|
self.assertIn('100%', text)
|
|
|
|
def test_progress_disabled_is_silent(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=1)
|
|
errors = StringIO()
|
|
with redirect_stderr(errors):
|
|
asyncio.run(generate(config, show_progress=False))
|
|
export(config, show_progress=False)
|
|
self.assertEqual(errors.getvalue(), '')
|
|
|
|
def test_empty_dataset_progress(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=0)
|
|
errors = StringIO()
|
|
with redirect_stderr(errors):
|
|
asyncio.run(generate(config, show_progress=True))
|
|
export(config, show_progress=True)
|
|
self.assertIn('0/0', errors.getvalue().replace(' ', ''))
|