Spaces:
Running
Running
File size: 2,047 Bytes
81e3ca2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | import tempfile
import unittest
from pathlib import Path
from ai_prof.deck_cache import DeckCache
from ai_prof.pdf_utils import Deck, Slide
class DeckCacheTests(unittest.TestCase):
def test_save_and_load_round_trip(self):
with tempfile.TemporaryDirectory() as tmp:
source = Path(tmp) / "source.png"
source.write_bytes(b"png")
cache = DeckCache(root=str(Path(tmp) / "cache"))
deck = Deck(
slides=[
Slide(index=0, image_path=str(source), text="First slide"),
]
)
cache.save(
"deck-key",
deck=deck,
readings={0: "TITLE: First\nCONCEPTS: caching"},
deck_index="1. First - caching",
metadata={"title": "Caching 101", "dpi": 150},
)
loaded = cache.load("deck-key")
self.assertIsNotNone(loaded)
self.assertEqual(loaded.deck.slides[0].text, "First slide")
self.assertEqual(
loaded.readings[0],
"TITLE: First\nCONCEPTS: caching",
)
self.assertEqual(loaded.deck_index, "1. First - caching")
self.assertTrue(Path(loaded.deck.slides[0].image_path).is_file())
self.assertEqual(
cache.list_decks()[0].title,
"Caching 101",
)
def test_key_changes_with_processing_settings(self):
with tempfile.TemporaryDirectory() as tmp:
pdf = Path(tmp) / "lecture.pdf"
pdf.write_bytes(b"same lecture")
cache = DeckCache(root=str(Path(tmp) / "cache"))
first = cache.key(str(pdf), dpi=150, vision_model="minicpm-v")
second = cache.key(str(pdf), dpi=200, vision_model="minicpm-v")
third = cache.key(str(pdf), dpi=150, vision_model="other-model")
self.assertNotEqual(first, second)
self.assertNotEqual(first, third)
if __name__ == "__main__":
unittest.main()
|