from __future__ import annotations import hashlib import tempfile import unittest from pathlib import Path from arr_storage.filesystem import FilesystemObjectBackend from arr_storage.store import ManagedObjectStore from arr_web.downloads import ( ArtifactDescriptor, ControlledProjectArtifactReader, ManagedObjectArtifactReader, RoutedArtifactReader, ) from company_reports.publishing import AtomicReportPublisher as CompanyPublisher from company_reports.repository import ReservedReport as CompanyReservation from monthly_reports.publishing import ( AtomicReportPublisher as MonthlyPublisher, OpenpyxlWorkbookBuilder, private_staging_directory, ) from monthly_reports.repository import ReservedReport as MonthlyReservation from tests.test_company_reports_publishing import ( FakeRepository as CompanyFakeRepository, empty_report, synthetic_built_workbook, ) from tests.test_monthly_reports_publishing import FakeRepository as MonthlyFakeRepository from tests.test_monthly_reports_xlsx import report_fixture XLSX_MIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" class ReportArtifactStorageTests(unittest.TestCase): def test_monthly_and_company_publications_upload_workbook_and_result_to_oss(self): with tempfile.TemporaryDirectory(prefix="report-artifact-oss-") as temporary: root = Path(temporary) / "project" root.mkdir() store = ManagedObjectStore( FilesystemObjectBackend(root / "objects", create=True) ) monthly_report = report_fixture() monthly_work = root / "monthly-work" monthly_work.mkdir() monthly_built = OpenpyxlWorkbookBuilder().build(monthly_report, monthly_work) monthly_outcome = MonthlyPublisher( root, root / "outputs" / "monthly_reports", object_store=store, ).publish( monthly_report, MonthlyReservation(44, 2), monthly_built, MonthlyFakeRepository(), monthly_work, ) self.assertEqual(monthly_outcome.artifact.storage_provider, "oss") self.assertEqual(monthly_outcome.result_json.storage_provider, "oss") self.assertEqual( store.inspect_committed( monthly_outcome.artifact.storage_key, monthly_report.filename, ).role, "monthly_report", ) self.assertEqual( store.inspect_committed( monthly_outcome.result_json.storage_key, monthly_outcome.result_json.original_filename, ).role, "result_json", ) company_report = empty_report() company_work = root / "company-work" company_work.mkdir() company_built = synthetic_built_workbook( company_work, company_report.filename, b"synthetic-company-xlsx", "a" * 64, ) company_outcome = CompanyPublisher( root, root / "outputs" / "company_reports", object_store=store, ).publish( company_report, CompanyReservation(45, 3, company_report.company), company_built, CompanyFakeRepository(), company_work, ) self.assertEqual(company_outcome.artifact.storage_provider, "oss") self.assertEqual(company_outcome.result_json.storage_provider, "oss") self.assertEqual( store.inspect_committed( company_outcome.artifact.storage_key, company_report.filename, ).role, "company_report", ) retry_work = root / "company-retry-work" retry_work.mkdir() retry_built = synthetic_built_workbook( retry_work, company_report.filename, b"synthetic-company-xlsx-rebuilt", "a" * 64, ) retry_outcome = CompanyPublisher( root, root / "outputs" / "company_reports", object_store=store, ).publish( company_report, CompanyReservation(45, 3, company_report.company), retry_built, CompanyFakeRepository(), retry_work, ) self.assertEqual(retry_outcome.artifact.storage_provider, "oss") self.assertEqual( retry_outcome.artifact.sha256, company_outcome.artifact.sha256, ) self.assertEqual(retry_outcome.result_json.storage_provider, "oss") def test_download_router_supports_oss_and_legacy_local_records(self): with tempfile.TemporaryDirectory(prefix="report-download-") as temporary: root = Path(temporary) source = root / "monthly.xlsx" source.write_bytes(b"monthly-oss-bytes") store = ManagedObjectStore( FilesystemObjectBackend(root / "objects", create=True) ) stored = store.upload_committed( job_id="monthly-download-001", attempt_no=1, role="monthly_report", source=source, original_filename="月报.xlsx", ) oss_descriptor = ArtifactDescriptor( "monthly_xlsx", "月报.xlsx", stored.object_key, stored.sha256, stored.byte_size, stored.mime_type, "oss", "arr-private", ) local_path = root / "outputs" / "legacy.xlsx" local_path.parent.mkdir() local_path.write_bytes(b"legacy-local-bytes") local_descriptor = ArtifactDescriptor( "monthly_xlsx", "legacy.xlsx", "outputs/legacy.xlsx", hashlib.sha256(local_path.read_bytes()).hexdigest(), local_path.stat().st_size, XLSX_MIME, ) reader = RoutedArtifactReader( daily_reader=None, local_reader=ControlledProjectArtifactReader(root), oss_reader=ManagedObjectArtifactReader(store), ) self.assertEqual(reader.read(oss_descriptor), source.read_bytes()) self.assertEqual(reader.read(local_descriptor), local_path.read_bytes()) if __name__ == "__main__": unittest.main()