Climate-Mood-Analysis/tests/test_nsrdb_polygon_metric_wrappers.py

from __future__ import annotations

import sys
import unittest
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))

import download_nsrdb_county_polygon_cloud_archives as cloud_download  # noqa: E402
import download_nsrdb_county_polygon_ghi_archives as ghi_download  # noqa: E402
import request_nsrdb_county_polygon_cloud_archives as cloud_request  # noqa: E402
import request_nsrdb_county_polygon_ghi_archives as ghi_request  # noqa: E402


class PolygonCloudWrapperTests(unittest.TestCase):
    def test_request_uses_cloud_artifact_label(self) -> None:
        args = cloud_request.polygon_request.parse_args(cloud_request.DEFAULT_ARGS)
        with TemporaryDirectory() as temp_dir:
            args.response_dir = Path(temp_dir)
            path = cloud_request.polygon_request.response_path(
                args,
                "01001",
                "goes-tmy",
                "tmy-2024",
            )

        self.assertEqual(
            path.name,
            "01001_goes-tmy_tmy-2024_cloud_response.json",
        )

    def test_request_reuses_legacy_cloud_response_name(self) -> None:
        args = cloud_request.polygon_request.parse_args(cloud_request.DEFAULT_ARGS)
        with TemporaryDirectory() as temp_dir:
            args.response_dir = Path(temp_dir)
            legacy_path = args.response_dir / "01001_goes-tmy_tmy-2024_ghi_response.json"
            legacy_path.touch()

            path = cloud_request.polygon_request.response_path(
                args,
                "01001",
                "goes-tmy",
                "tmy-2024",
            )

        self.assertEqual(path, legacy_path)

    def test_request_user_arguments_override_cloud_defaults(self) -> None:
        args = cloud_request.polygon_request.parse_args(
            [
                *cloud_request.DEFAULT_ARGS,
                "--attributes",
                "ghi",
                "--requests-csv",
                "custom.csv",
            ]
        )

        self.assertEqual(args.attributes, "ghi")
        self.assertEqual(args.requests_csv, Path("custom.csv"))

    def test_request_wrapper_passes_cloud_defaults_and_user_arguments(self) -> None:
        with (
            patch.object(sys, "argv", ["cloud-request", "--limit", "12"]),
            patch.object(cloud_request.polygon_request, "main") as shared_main,
        ):
            cloud_request.main()

        shared_main.assert_called_once_with(
            [*cloud_request.DEFAULT_ARGS, "--limit", "12"],
            description=cloud_request.__doc__,
        )

    def test_download_wrapper_passes_cloud_defaults_and_user_arguments(self) -> None:
        with (
            patch.object(sys, "argv", ["cloud-download", "--overwrite"]),
            patch.object(cloud_download.polygon_download, "main") as shared_main,
        ):
            cloud_download.main()

        shared_main.assert_called_once_with(
            [*cloud_download.DEFAULT_ARGS, "--overwrite"],
            description=cloud_download.__doc__,
        )

    def test_download_user_arguments_override_cloud_defaults(self) -> None:
        args = cloud_download.polygon_download.parse_args(
            [
                *cloud_download.DEFAULT_ARGS,
                "--output-dir",
                "custom-archives",
            ]
        )

        self.assertEqual(args.output_dir, Path("custom-archives"))


class PolygonGhiWrapperTests(unittest.TestCase):
    def test_request_uses_ghi_artifact_label(self) -> None:
        args = ghi_request.polygon_request.parse_args(ghi_request.DEFAULT_ARGS)
        with TemporaryDirectory() as temp_dir:
            args.response_dir = Path(temp_dir)
            path = ghi_request.polygon_request.response_path(
                args,
                "01001",
                "goes-tmy",
                "tmy-2024",
            )

        self.assertEqual(
            path.name,
            "01001_goes-tmy_tmy-2024_ghi_response.json",
        )

    def test_request_wrapper_passes_ghi_defaults_and_user_arguments(self) -> None:
        with (
            patch.object(sys, "argv", ["ghi-request", "--limit", "12"]),
            patch.object(ghi_request.polygon_request, "main") as shared_main,
        ):
            ghi_request.main()

        shared_main.assert_called_once_with(
            [*ghi_request.DEFAULT_ARGS, "--limit", "12"],
            description=ghi_request.__doc__,
        )

    def test_request_user_arguments_override_ghi_defaults(self) -> None:
        args = ghi_request.polygon_request.parse_args(
            [
                *ghi_request.DEFAULT_ARGS,
                "--attributes",
                "ghi,clearsky_ghi",
                "--requests-csv",
                "custom.csv",
            ]
        )

        self.assertEqual(args.attributes, "ghi,clearsky_ghi")
        self.assertEqual(args.requests_csv, Path("custom.csv"))

    def test_download_wrapper_passes_ghi_defaults_and_user_arguments(self) -> None:
        with (
            patch.object(sys, "argv", ["ghi-download", "--overwrite"]),
            patch.object(ghi_download.polygon_download, "main") as shared_main,
        ):
            ghi_download.main()

        shared_main.assert_called_once_with(
            [*ghi_download.DEFAULT_ARGS, "--overwrite"],
            description=ghi_download.__doc__,
        )

    def test_download_user_arguments_override_ghi_defaults(self) -> None:
        args = ghi_download.polygon_download.parse_args(
            [
                *ghi_download.DEFAULT_ARGS,
                "--output-dir",
                "custom-archives",
            ]
        )

        self.assertEqual(args.output_dir, Path("custom-archives"))


if __name__ == "__main__":
    unittest.main()