130 lines
4.8 KiB
Python
130 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from openlane_scraper.core.config import OpenLaneConfig
|
|
from openlane_scraper.openlane.checkpoint import OpenLaneCheckpoint, OpenLaneCheckpointStore
|
|
from openlane_scraper.openlane.client import OpenLaneClient
|
|
from openlane_scraper.openlane.writer import OpenLaneResultWriter
|
|
|
|
|
|
class TestOpenLaneCheckpointStore(unittest.TestCase):
|
|
def test_save_and_load_checkpoint(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = Path(tmp_dir) / "checkpoint.json"
|
|
store = OpenLaneCheckpointStore(path)
|
|
checkpoint = OpenLaneCheckpoint(
|
|
max_pages=512,
|
|
completed_pages=[1, 2, 3],
|
|
failed_pages=[7],
|
|
total_records=99,
|
|
last_saved_at="2026-04-20T00:00:00+00:00",
|
|
)
|
|
|
|
store.save(checkpoint)
|
|
loaded = store.load(max_pages=512)
|
|
|
|
self.assertEqual(loaded.max_pages, 512)
|
|
self.assertEqual(loaded.completed_pages, [1, 2, 3])
|
|
self.assertEqual(loaded.failed_pages, [7])
|
|
self.assertEqual(loaded.total_records, 99)
|
|
self.assertEqual(loaded.last_saved_at, "2026-04-20T00:00:00+00:00")
|
|
|
|
|
|
class TestOpenLaneWriter(unittest.TestCase):
|
|
def test_append_and_finalize_outputs(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
jsonl_path = Path(tmp_dir) / "openlane.jsonl"
|
|
aggregated_path = Path(tmp_dir) / "openlane_aggregated.json"
|
|
writer = OpenLaneResultWriter(jsonl_path, aggregated_path)
|
|
|
|
written = writer.append_page(1, [{"id": 1}, {"id": 2}])
|
|
self.assertEqual(written, 2)
|
|
|
|
summary = writer.finalize(max_pages=10, completed_pages=[1], failed_pages=[2])
|
|
|
|
self.assertEqual(summary["total_records"], 2)
|
|
self.assertEqual(summary["total_pages_completed"], 1)
|
|
self.assertEqual(summary["total_pages_failed"], 1)
|
|
self.assertTrue(aggregated_path.exists())
|
|
|
|
payload = json.loads(aggregated_path.read_text(encoding="utf-8"))
|
|
self.assertEqual(payload["items"][0]["page"], 1)
|
|
self.assertEqual(payload["items"][0]["record"]["id"], 1)
|
|
|
|
|
|
class TestOpenLaneConfig(unittest.TestCase):
|
|
def test_retry_schedule_defaults(self) -> None:
|
|
config = OpenLaneConfig()
|
|
self.assertGreaterEqual(len(config.retry_schedule_seconds), 1)
|
|
self.assertEqual(tuple(float(x) for x in config.retry_schedule_seconds), config.retry_schedule_seconds)
|
|
|
|
def test_storage_state_defaults(self) -> None:
|
|
config = OpenLaneConfig()
|
|
self.assertTrue(config.storage_state_file)
|
|
self.assertIn("artifacts/openlane", config.storage_state_file)
|
|
self.assertIsInstance(config.persist_storage_state, bool)
|
|
|
|
|
|
class TestOpenLaneClient(unittest.TestCase):
|
|
def test_extract_records_from_common_shapes(self) -> None:
|
|
self.assertEqual(
|
|
OpenLaneClient._extract_records({"results": [{"id": 1}, {"id": 2}]}),
|
|
[{"id": 1}, {"id": 2}],
|
|
)
|
|
self.assertEqual(
|
|
OpenLaneClient._extract_records({"data": {"items": [{"id": 3}]}}),
|
|
[{"id": 3}],
|
|
)
|
|
self.assertEqual(
|
|
OpenLaneClient._extract_records({"vehicles": [{"id": 4}]}),
|
|
[{"id": 4}],
|
|
)
|
|
self.assertEqual(OpenLaneClient._extract_records({"data": {"foo": "bar"}}), [])
|
|
|
|
def test_extract_vehicle_images_from_common_shapes(self) -> None:
|
|
payload = {
|
|
"vehicle_images": [
|
|
{"url": "https://img/1.png", "low_resolution_url": "https://img/1_small.png"},
|
|
]
|
|
}
|
|
self.assertEqual(
|
|
OpenLaneClient._extract_vehicle_images(payload),
|
|
[{"url": "https://img/1.png", "low_resolution_url": "https://img/1_small.png"}],
|
|
)
|
|
|
|
payload_nested = {
|
|
"data": {
|
|
"images": [
|
|
{"url": "https://img/2.png"},
|
|
]
|
|
}
|
|
}
|
|
self.assertEqual(
|
|
OpenLaneClient._extract_vehicle_images(payload_nested),
|
|
[{"url": "https://img/2.png"}],
|
|
)
|
|
|
|
self.assertEqual(OpenLaneClient._extract_vehicle_images({"data": {"foo": 1}}), [])
|
|
|
|
def test_extract_vehicle_images_batch_helper_normalization(self) -> None:
|
|
# Косвенно проверяем нормализацию входа под батчевый метод
|
|
ids = [" 123 ", "123", 456, "", None]
|
|
normalized = []
|
|
seen = set()
|
|
for v in ids:
|
|
s = str(v).strip()
|
|
if not s or s in seen:
|
|
continue
|
|
seen.add(s)
|
|
normalized.append(s)
|
|
|
|
self.assertEqual(normalized, ["123", "456", "None"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|