Initial Climate Mood Analysis project

2026-06-11 14:50:33 -04:00
commit 18e099ce84
34 changed files with 18807 additions and 0 deletions
@@ -0,0 +1,163 @@
+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()