예제 소스

test_chart_png.py

import tempfile
import unittest
from pathlib import Path

from lib.chart_png import (
    CHART_BUCKET_SIZE,
    chart_png_path,
    existing_chart_bars,
    migrate_flat_chart_dir,
)


class ChartPngPathTest(unittest.TestCase):
    def test_bucket_is_under_10000(self):
        self.assertEqual(CHART_BUCKET_SIZE, 10_000)
        root = Path("/tmp/chart")
        self.assertEqual(chart_png_path(root, 300), root / "0000" / "0000300.png")
        self.assertEqual(chart_png_path(root, 9999), root / "0000" / "0009999.png")
        self.assertEqual(chart_png_path(root, 10000), root / "0001" / "0010000.png")
        self.assertEqual(chart_png_path(root, 242228), root / "0024" / "0242228.png")


class ChartPngMigrateTest(unittest.TestCase):
    def test_moves_flat_pngs_into_buckets(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            (root / "0000300.png").write_bytes(b"a" * 10)
            (root / "0010000.png").write_bytes(b"b" * 10)
            moved, skipped = migrate_flat_chart_dir(root)
            self.assertEqual(moved, 2)
            self.assertEqual(skipped, 0)
            self.assertFalse((root / "0000300.png").exists())
            self.assertTrue((root / "0000" / "0000300.png").is_file())
            self.assertTrue((root / "0001" / "0010000.png").is_file())
            self.assertEqual(existing_chart_bars(root), {300, 10000})