111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
import base64
|
|
import tempfile
|
|
import unittest
|
|
import urllib.error
|
|
from email.message import Message
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from ingestion.minjust_cbd import CbdClient, SyncResult, progress_line, sync_archive
|
|
|
|
|
|
class MinjustCbdTest(unittest.TestCase):
|
|
def test_archives_document_and_resumes(self):
|
|
image = base64.b64encode(b"image").decode()
|
|
document = {
|
|
"Code": 1,
|
|
"Name": {"Rus": "Закон", "Kyr": "Мыйзам"},
|
|
"Editions": [
|
|
{
|
|
"Code": 10,
|
|
"Name": {"Rus": "01.01.2026", "Kyr": "01.01.2026"},
|
|
"Data": {"Rus": "<p>Закон</p>", "Kyr": "<p>Мыйзам</p>"},
|
|
"Images": [
|
|
{
|
|
"Lang": "Russian",
|
|
"Name": "image001.jpg",
|
|
"Data": image,
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
client = Mock()
|
|
client.total_documents = 1
|
|
client.document_codes.return_value = [1]
|
|
client.document.return_value = document
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
output = Path(directory)
|
|
first = sync_archive(output, client)
|
|
second = sync_archive(output, client)
|
|
refreshed = sync_archive(output, client, refresh=True)
|
|
|
|
self.assertEqual(first.downloaded, 1)
|
|
self.assertEqual(second.skipped, 1)
|
|
self.assertEqual(refreshed.downloaded, 1)
|
|
self.assertEqual(client.document.call_count, 2)
|
|
self.assertEqual(
|
|
(output / "documents/1/editions/10/ru.html").read_text(),
|
|
"<p>Закон</p>",
|
|
)
|
|
self.assertEqual(
|
|
(output / "documents/1/editions/10/images/ru/image001.jpg").read_bytes(),
|
|
b"image",
|
|
)
|
|
|
|
def test_uses_stable_list_id_for_next_page(self):
|
|
client = CbdClient(requests_per_second=1000)
|
|
client.request_json = Mock(
|
|
side_effect=[
|
|
{
|
|
"Id": "list-id",
|
|
"TotalCount": 3,
|
|
"Documents": [{"Code": 1}, {"Code": 2}],
|
|
},
|
|
{
|
|
"Id": "list-id",
|
|
"TotalCount": 3,
|
|
"Documents": [{"Code": 3}],
|
|
},
|
|
]
|
|
)
|
|
|
|
self.assertEqual(list(client.document_codes(page_size=2)), [1, 2, 3])
|
|
self.assertEqual(
|
|
client.request_json.call_args_list[1].args[0], "GetDocumentListById"
|
|
)
|
|
|
|
def test_formats_progress_with_rate_and_eta(self):
|
|
line = progress_line(
|
|
SyncResult(discovered=50, downloaded=48, skipped=1, failed=1),
|
|
total=100,
|
|
elapsed=25,
|
|
)
|
|
|
|
self.assertIn("50.00%", line)
|
|
self.assertIn("rate=2.00/s", line)
|
|
self.assertIn("ETA=00:00:25", line)
|
|
|
|
@patch("ingestion.minjust_cbd.time.sleep")
|
|
@patch("ingestion.minjust_cbd.urllib.request.urlopen")
|
|
def test_respects_retry_after_on_rate_limit(self, urlopen, sleep):
|
|
headers = Message()
|
|
headers["Retry-After"] = "7"
|
|
response = Mock()
|
|
response.read.return_value = b"{}"
|
|
response.__enter__ = Mock(return_value=response)
|
|
response.__exit__ = Mock(return_value=False)
|
|
urlopen.side_effect = [
|
|
urllib.error.HTTPError("https://example.test", 429, "", headers, None),
|
|
response,
|
|
]
|
|
|
|
CbdClient(requests_per_second=1000, retries=2).request_json("Example")
|
|
|
|
self.assertIn(((7.0,), {}), [(call.args, call.kwargs) for call in sleep.call_args_list])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|