crypto_monitor/tests/test_journal_images_lib.py

"""journal_images_lib / journal_upload_api_lib 单元测试."""
import json
import os
import tempfile
import unittest
from io import BytesIO

from lib.instance.journal_images_lib import (
    JOURNAL_UPLOAD_TFS,
    collect_journal_slot_images,
    enrich_journal_api_item,
    images_json_dumps,
    is_valid_preuploaded_journal_file,
    journal_image_paths,
    journal_upload_field_name,
    normalize_journal_draft_id,
    parse_images_json,
    primary_journal_image,
    save_journal_slot_uploads,
    uploaded_screenshot_field_name,
)
from lib.instance.journal_upload_api_lib import handle_journal_upload_slot


class _FakeFile:
    def __init__(self, filename: str, data: bytes):
        self.filename = filename
        self._data = data

    def save(self, path: str) -> None:
        with open(path, "wb") as f:
            f.write(self._data)


class _FakeFiles:
    def __init__(self, mapping):
        self._mapping = mapping

    def get(self, key):
        return self._mapping.get(key)


class _FakeForm:
    def __init__(self, mapping):
        self._mapping = mapping

    def get(self, key, default=None):
        return self._mapping.get(key, default)


class _FakeRequest:
    def __init__(self, form=None, files=None):
        self.form = form
        self.files = files


class JournalImagesLibTest(unittest.TestCase):
    def test_field_names(self):
        self.assertEqual(journal_upload_field_name("5m"), "screenshot_5m")
        self.assertEqual(uploaded_screenshot_field_name("5m"), "uploaded_screenshot_5m")

    def test_normalize_draft_id(self):
        good = "a" * 32
        self.assertEqual(normalize_journal_draft_id(good), good)
        self.assertIsNone(normalize_journal_draft_id("bad"))

    def test_save_slot_uploads_partial(self):
        with tempfile.TemporaryDirectory() as tmp:
            files = _FakeFiles(
                {
                    "screenshot_5m": _FakeFile("a.png", b"png5"),
                    "screenshot_1h": _FakeFile("b.jpg", b"jpg1"),
                }
            )
            saved = save_journal_slot_uploads(
                files,
                "abc123" + "0" * 26,
                tmp,
                secure_filename_fn=lambda x: x,
            )
            self.assertEqual(len(saved), 2)
            self.assertEqual(saved[0]["tf"], "5m")
            self.assertTrue(os.path.isfile(os.path.join(tmp, saved[0]["file"])))
            self.assertEqual(saved[1]["tf"], "1h")

    def test_collect_preuploaded(self):
        entry_id = "abc123" + "0" * 26
        fname = f"journal_{entry_id}_5m.png"
        with tempfile.TemporaryDirectory() as tmp:
            with open(os.path.join(tmp, fname), "wb") as f:
                f.write(b"x")
            form = _FakeForm({uploaded_screenshot_field_name("5m"): fname})
            saved = collect_journal_slot_images(
                form,
                _FakeFiles({}),
                entry_id,
                tmp,
                secure_filename_fn=lambda x: x,
            )
            self.assertEqual(saved, [{"tf": "5m", "file": fname}])

    def test_is_valid_preuploaded_journal_file(self):
        entry_id = "abc123" + "0" * 26
        fname = f"journal_{entry_id}_5m.png"
        self.assertTrue(is_valid_preuploaded_journal_file(fname, entry_id, "5m"))
        self.assertFalse(is_valid_preuploaded_journal_file("../evil.png", entry_id, "5m"))
        self.assertFalse(is_valid_preuploaded_journal_file(fname, "b" * 32, "5m"))

    def test_parse_and_enrich(self):
        raw = images_json_dumps([{"tf": "5m", "file": "journal_x_5m.png"}])
        item = enrich_journal_api_item({"images_json": raw, "image": "legacy.png"})
        self.assertEqual(len(item["images"]), 1)
        self.assertEqual(item["images"][0]["tf"], "5m")

        legacy = enrich_journal_api_item({"image": "only.png"})
        self.assertEqual(legacy["images"][0]["file"], "only.png")

    def test_journal_image_paths_dedupe(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "same.png")
            with open(path, "wb") as f:
                f.write(b"x")
            row = {
                "image": "same.png",
                "images_json": json.dumps([{"tf": "5m", "file": "same.png"}]),
            }
            paths = journal_image_paths(row, tmp)
            self.assertEqual(len(paths), 1)

    def test_primary_journal_image(self):
        self.assertEqual(
            primary_journal_image([{"tf": "5m", "file": "a.png"}]),
            "a.png",
        )
        self.assertIsNone(primary_journal_image([]))

    def test_handle_journal_upload_slot(self):
        entry_id = "abc123" + "0" * 26
        with tempfile.TemporaryDirectory() as tmp:
            req = _FakeRequest(
                form=_FakeForm({"journal_draft_id": entry_id, "tf": "5m"}),
                files=_FakeFiles({"file": _FakeFile("local.png", b"data")}),
            )
            payload, code = handle_journal_upload_slot(
                req,
                upload_folder=tmp,
                secure_filename_fn=lambda x: x,
            )
            self.assertEqual(code, 200)
            self.assertTrue(payload["ok"])
            self.assertEqual(payload["tf"], "5m")
            self.assertTrue(os.path.isfile(os.path.join(tmp, payload["file"])))


if __name__ == "__main__":
    unittest.main()