from __future__ import annotations import sys import unittest import urllib.error from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import patch SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts" sys.path.insert(0, str(SCRIPTS_DIR)) from download_nsrdb_county_polygon_archives import ( # noqa: E402 ArchiveDownloadEvent, ArchiveDownloadState, ArchiveDownloadStateMachine, InvalidArchiveDownloadTransition, PolygonArchiveDownloadError, download_file, output_name_for_response, ) class FakeResponse: def __init__(self, content: bytes) -> None: self.content = content self.offset = 0 def __enter__(self) -> FakeResponse: return self def __exit__(self, *args: object) -> None: return None def read(self, size: int) -> bytes: chunk = self.content[self.offset : self.offset + size] self.offset += len(chunk) return chunk class ArchiveDownloadStateMachineTests(unittest.TestCase): def test_legacy_ghi_response_name_keeps_existing_archive_name(self) -> None: response_path = Path("01001_goes-tmy_tmy-2024_ghi_response.json") self.assertEqual( output_name_for_response(response_path), "01001_goes-tmy_tmy-2024_ghi.zip", ) def test_successful_download_lifecycle(self) -> None: machine = ArchiveDownloadStateMachine("01001_response.json") machine.transition(ArchiveDownloadEvent.START) with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "01001.zip" with patch( "download_nsrdb_county_polygon_archives.urllib.request.urlopen", return_value=FakeResponse(b"archive"), ): state = download_file( "https://example.com/01001.zip", output_path, timeout=30, overwrite=False, machine=machine, ) self.assertEqual(output_path.read_bytes(), b"archive") self.assertEqual(state, ArchiveDownloadState.DOWNLOADED) self.assertEqual( machine.history, [ ArchiveDownloadState.QUEUED, ArchiveDownloadState.CHECKING, ArchiveDownloadState.DOWNLOADING, ArchiveDownloadState.DOWNLOADED, ], ) def test_existing_archive_is_skipped(self) -> None: machine = ArchiveDownloadStateMachine("01001_response.json") machine.transition(ArchiveDownloadEvent.START) with TemporaryDirectory() as temp_dir: output_path = Path(temp_dir) / "01001.zip" output_path.write_bytes(b"existing") state = download_file( "https://example.com/01001.zip", output_path, timeout=30, overwrite=False, machine=machine, ) self.assertEqual(state, ArchiveDownloadState.SKIPPED) def test_pending_s3_archive_enters_pending_state(self) -> None: machine = ArchiveDownloadStateMachine("01001_response.json") machine.transition(ArchiveDownloadEvent.START) error = urllib.error.HTTPError( "https://bucket.s3.amazonaws.com/01001.zip", 403, "Forbidden", {}, None, ) with TemporaryDirectory() as temp_dir: with patch( "download_nsrdb_county_polygon_archives.urllib.request.urlopen", side_effect=error, ): state = download_file( error.url, Path(temp_dir) / "01001.zip", timeout=30, overwrite=False, machine=machine, ) self.assertEqual(state, ArchiveDownloadState.PENDING) def test_http_error_enters_failed_state(self) -> None: machine = ArchiveDownloadStateMachine("01001_response.json") machine.transition(ArchiveDownloadEvent.START) error = urllib.error.HTTPError( "https://example.com/01001.zip", 500, "Server Error", {}, None, ) with TemporaryDirectory() as temp_dir: with patch( "download_nsrdb_county_polygon_archives.urllib.request.urlopen", side_effect=error, ): with self.assertRaises(PolygonArchiveDownloadError): download_file( error.url, Path(temp_dir) / "01001.zip", timeout=30, overwrite=False, machine=machine, ) self.assertEqual(machine.state, ArchiveDownloadState.FAILED) def test_invalid_transition_is_rejected(self) -> None: machine = ArchiveDownloadStateMachine("01001_response.json") with self.assertRaisesRegex( InvalidArchiveDownloadTransition, "download_succeeded while queued", ): machine.transition(ArchiveDownloadEvent.DOWNLOAD_SUCCEEDED) if __name__ == "__main__": unittest.main()