69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
import json
|
|
|
|
from haiku.rag.store.compression import (
|
|
compress_docling_split,
|
|
compress_json,
|
|
decompress_json,
|
|
)
|
|
|
|
|
|
class TestJsonCompression:
|
|
def test_compress_json_roundtrip(self):
|
|
json_str = '{"key": "value", "number": 42, "nested": {"a": 1}}'
|
|
compressed = compress_json(json_str)
|
|
decompressed = decompress_json(compressed)
|
|
assert decompressed == json_str
|
|
|
|
def test_compress_json_with_unicode(self):
|
|
json_str = '{"message": "Hello, 世界! 🌍"}'
|
|
compressed = compress_json(json_str)
|
|
decompressed = decompress_json(compressed)
|
|
assert decompressed == json_str
|
|
|
|
def test_compress_json_produces_zstd(self):
|
|
compressed = compress_json('{"test": true}')
|
|
assert compressed[:4] == b"\x28\xb5\x2f\xfd"
|
|
|
|
|
|
class TestDoclingCompressionSplit:
|
|
def test_split_with_pages(self):
|
|
data = {
|
|
"name": "test_doc",
|
|
"texts": [{"text": "hello"}],
|
|
"pages": {"1": {"image": "base64data"}, "2": {"image": "more"}},
|
|
}
|
|
json_str = json.dumps(data)
|
|
structure_bytes, pages_bytes = compress_docling_split(json_str)
|
|
|
|
assert structure_bytes is not None
|
|
assert pages_bytes is not None
|
|
|
|
# Structure should not contain pages
|
|
structure = json.loads(decompress_json(structure_bytes))
|
|
assert "pages" not in structure
|
|
assert structure["name"] == "test_doc"
|
|
assert structure["texts"] == [{"text": "hello"}]
|
|
|
|
# Pages should contain only pages
|
|
pages = json.loads(decompress_json(pages_bytes))
|
|
assert "1" in pages
|
|
assert "2" in pages
|
|
|
|
def test_split_without_pages(self):
|
|
data = {"name": "test_doc", "texts": []}
|
|
json_str = json.dumps(data)
|
|
structure_bytes, pages_bytes = compress_docling_split(json_str)
|
|
|
|
assert structure_bytes is not None
|
|
assert pages_bytes is None
|
|
|
|
structure = json.loads(decompress_json(structure_bytes))
|
|
assert structure["name"] == "test_doc"
|
|
|
|
def test_split_with_empty_pages(self):
|
|
data = {"name": "test_doc", "texts": [], "pages": {}}
|
|
json_str = json.dumps(data)
|
|
structure_bytes, pages_bytes = compress_docling_split(json_str)
|
|
|
|
assert structure_bytes is not None
|
|
assert pages_bytes is None
|