Climate-Mood-Analysis/tests/test_nsrdb_polygon_download_state_machine.py

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()