Initial Climate Mood Analysis project
This commit is contained in:
Justin Fisher
@@ -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()
|
||||
@@ -0,0 +1,170 @@
|
||||
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()
|
||||
@@ -0,0 +1,355 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
import csv
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from unittest.mock import call, patch
|
||||
|
||||
|
||||
SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts"
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
from request_nsrdb_county_polygon_archives import ( # noqa: E402
|
||||
ArchiveQueueMonitor,
|
||||
CountyRequestEvent,
|
||||
CountyRequestState,
|
||||
CountyRequestStateMachine,
|
||||
EXCEPTIONAL_POLYGON_WAIT,
|
||||
InvalidCountyRequestTransition,
|
||||
LARGE_POLYGON_WAIT,
|
||||
LocalQueueCapacityError,
|
||||
MEDIUM_POLYGON_WAIT,
|
||||
SMALL_POLYGON_WAIT,
|
||||
)
|
||||
|
||||
|
||||
class CountyRequestStateMachineTests(unittest.TestCase):
|
||||
def test_successful_request_lifecycle(self) -> None:
|
||||
machine = CountyRequestStateMachine("01001 Autauga County, AL")
|
||||
|
||||
machine.transition(CountyRequestEvent.START)
|
||||
machine.transition(CountyRequestEvent.REQUEST_SUCCEEDED)
|
||||
|
||||
self.assertEqual(machine.state, CountyRequestState.REQUESTED)
|
||||
self.assertEqual(
|
||||
machine.history,
|
||||
[
|
||||
CountyRequestState.PENDING,
|
||||
CountyRequestState.SUBMITTING,
|
||||
CountyRequestState.REQUESTED,
|
||||
],
|
||||
)
|
||||
|
||||
def test_retry_returns_to_submitting(self) -> None:
|
||||
machine = CountyRequestStateMachine("01001 Autauga County, AL")
|
||||
|
||||
machine.transition(CountyRequestEvent.START)
|
||||
machine.transition(CountyRequestEvent.RETRY_REQUIRED)
|
||||
machine.transition(CountyRequestEvent.RETRY_STARTED)
|
||||
|
||||
self.assertEqual(machine.state, CountyRequestState.SUBMITTING)
|
||||
|
||||
def test_existing_request_is_skipped(self) -> None:
|
||||
machine = CountyRequestStateMachine("01001 Autauga County, AL")
|
||||
|
||||
machine.transition(CountyRequestEvent.EXISTING_REQUEST_FOUND)
|
||||
|
||||
self.assertEqual(machine.state, CountyRequestState.SKIPPED)
|
||||
|
||||
def test_invalid_transition_is_rejected(self) -> None:
|
||||
machine = CountyRequestStateMachine("01001 Autauga County, AL")
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
InvalidCountyRequestTransition,
|
||||
"request_succeeded while pending",
|
||||
):
|
||||
machine.transition(CountyRequestEvent.REQUEST_SUCCEEDED)
|
||||
|
||||
|
||||
class ArchiveQueueMonitorTests(unittest.TestCase):
|
||||
def make_monitor(self) -> ArchiveQueueMonitor:
|
||||
temp_dir = TemporaryDirectory()
|
||||
self.addCleanup(temp_dir.cleanup)
|
||||
return ArchiveQueueMonitor(Path(temp_dir.name), timeout=1, workers=1)
|
||||
|
||||
def test_dynamic_wait_uses_largest_tracked_polygon(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
|
||||
for site_count, expected_wait in (
|
||||
(110, SMALL_POLYGON_WAIT),
|
||||
(111, MEDIUM_POLYGON_WAIT),
|
||||
(250, MEDIUM_POLYGON_WAIT),
|
||||
(251, LARGE_POLYGON_WAIT),
|
||||
(550, LARGE_POLYGON_WAIT),
|
||||
(551, EXCEPTIONAL_POLYGON_WAIT),
|
||||
):
|
||||
monitor.tracked_site_counts = {"https://example.com/job.zip": site_count}
|
||||
self.assertEqual(monitor.dynamic_wait_seconds(), expected_wait)
|
||||
|
||||
def test_downloaded_archive_is_not_reloaded_as_pending(self) -> None:
|
||||
temp_dir = TemporaryDirectory()
|
||||
self.addCleanup(temp_dir.cleanup)
|
||||
manifest_dir = Path(temp_dir.name)
|
||||
response_dir = manifest_dir / "polygon_cloud_request_responses"
|
||||
archive_dir = manifest_dir / "polygon_cloud_archives"
|
||||
response_dir.mkdir()
|
||||
archive_dir.mkdir()
|
||||
response_path = response_dir / "01001_goes-tmy_tmy-2024_cloud_response.json"
|
||||
archive_path = archive_dir / "01001_goes-tmy_tmy-2024_cloud.zip"
|
||||
response_path.write_text("{}", encoding="utf-8")
|
||||
archive_path.write_bytes(b"downloaded")
|
||||
|
||||
manifest_path = manifest_dir / "county_polygon_cloud_request_manifest.csv"
|
||||
with manifest_path.open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(
|
||||
handle,
|
||||
fieldnames=[
|
||||
"county_fips",
|
||||
"site_count",
|
||||
"download_url",
|
||||
"response_json",
|
||||
"submitted_at_utc",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
writer.writerow(
|
||||
{
|
||||
"county_fips": "01001",
|
||||
"site_count": "100",
|
||||
"download_url": "https://example.amazonaws.com/job.zip",
|
||||
"response_json": str(response_path),
|
||||
"submitted_at_utc": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
monitor = ArchiveQueueMonitor(manifest_dir, timeout=1, workers=1)
|
||||
|
||||
with patch(
|
||||
"request_nsrdb_county_polygon_archives.archive_download_status",
|
||||
return_value="pending",
|
||||
) as status_mock:
|
||||
counts = monitor.summarize()
|
||||
|
||||
self.assertEqual(counts["pending"], 0)
|
||||
self.assertEqual(counts["total"], 0)
|
||||
status_mock.assert_not_called()
|
||||
|
||||
def test_historical_s3_403_does_not_consume_a_queue_slot(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
url = "https://example.amazonaws.com/historical.zip"
|
||||
monitor.jobs[url] = "01001"
|
||||
|
||||
with patch(
|
||||
"request_nsrdb_county_polygon_archives.archive_download_status",
|
||||
return_value="pending",
|
||||
):
|
||||
counts = monitor.summarize()
|
||||
|
||||
self.assertEqual(counts["pending"], 0)
|
||||
self.assertEqual(counts["ambiguous_403"], 1)
|
||||
|
||||
def test_current_process_s3_403_remains_pending(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
url = "https://example.amazonaws.com/current.zip"
|
||||
monitor.record_submission(
|
||||
{
|
||||
"county_fips": "01001",
|
||||
"tile_id": "",
|
||||
"site_count": "100",
|
||||
"download_url": url,
|
||||
}
|
||||
)
|
||||
|
||||
with patch(
|
||||
"request_nsrdb_county_polygon_archives.archive_download_status",
|
||||
return_value="pending",
|
||||
):
|
||||
counts = monitor.summarize()
|
||||
|
||||
self.assertEqual(counts["pending"], 1)
|
||||
self.assertEqual(counts["ambiguous_403"], 0)
|
||||
|
||||
def test_site_count_expires_after_recent_download_window(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
url = "https://example.com/job.zip"
|
||||
monitor.jobs[url] = "01001"
|
||||
monitor.pending_urls.add(url)
|
||||
monitor.tracked_site_counts[url] = 551
|
||||
|
||||
with (
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.archive_download_status",
|
||||
return_value="ready",
|
||||
),
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.time.monotonic",
|
||||
side_effect=[100.0, 100.0, 401.0],
|
||||
),
|
||||
):
|
||||
first_counts = monitor.summarize()
|
||||
second_counts = monitor.summarize()
|
||||
|
||||
self.assertEqual(first_counts["recently_downloaded"], 1)
|
||||
self.assertEqual(second_counts["recently_downloaded"], 0)
|
||||
self.assertNotIn(url, monitor.tracked_site_counts)
|
||||
|
||||
def test_fifteen_visible_slots_bypass_pacing_until_queue_is_full(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
monitor.paced_submissions = 1
|
||||
counts_with_fifteen_slots = {
|
||||
"pending": 5,
|
||||
"recently_downloaded": 2,
|
||||
"unobserved": 0,
|
||||
"unknown": 0,
|
||||
"ready": 0,
|
||||
"total": 7,
|
||||
}
|
||||
counts_with_fourteen_slots = {
|
||||
**counts_with_fifteen_slots,
|
||||
"pending": 6,
|
||||
}
|
||||
full_queue_counts = {
|
||||
**counts_with_fifteen_slots,
|
||||
"pending": 20,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
monitor,
|
||||
"summarize",
|
||||
side_effect=[
|
||||
counts_with_fifteen_slots,
|
||||
counts_with_fourteen_slots,
|
||||
full_queue_counts,
|
||||
],
|
||||
),
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.wait_for_queue",
|
||||
) as wait_for_queue_mock,
|
||||
):
|
||||
monitor.wait_for_capacity("first", retries=0)
|
||||
monitor.wait_for_capacity("second", retries=0)
|
||||
with self.assertRaisesRegex(LocalQueueCapacityError, "20 possible"):
|
||||
monitor.wait_for_capacity("third", retries=0)
|
||||
|
||||
wait_for_queue_mock.assert_not_called()
|
||||
self.assertFalse(monitor.filling_open_queue)
|
||||
|
||||
def test_fourteen_visible_slots_still_use_recent_download_pacing(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
monitor.paced_submissions = 1
|
||||
counts = {
|
||||
"pending": 6,
|
||||
"recently_downloaded": 2,
|
||||
"unobserved": 0,
|
||||
"unknown": 0,
|
||||
"ready": 0,
|
||||
"total": 8,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(monitor, "summarize", return_value=counts),
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.wait_for_queue",
|
||||
) as wait_for_queue_mock,
|
||||
):
|
||||
monitor.wait_for_capacity("next", retries=0)
|
||||
|
||||
wait_for_queue_mock.assert_called_once_with(
|
||||
SMALL_POLYGON_WAIT,
|
||||
"Pacing wait complete; refreshing archive status.",
|
||||
)
|
||||
|
||||
def test_six_minute_wait_checks_queue_after_three_minutes(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
counts = {
|
||||
"pending": 6,
|
||||
"recently_downloaded": 0,
|
||||
"unobserved": 0,
|
||||
"unknown": 0,
|
||||
"ready": 0,
|
||||
"total": 6,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(monitor, "summarize", return_value=counts) as summarize_mock,
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.time.sleep",
|
||||
) as sleep_mock,
|
||||
):
|
||||
entered_fill_mode = monitor.wait_for_dynamic_queue(
|
||||
LARGE_POLYGON_WAIT,
|
||||
"Queue wait complete.",
|
||||
)
|
||||
|
||||
self.assertFalse(entered_fill_mode)
|
||||
self.assertEqual(sleep_mock.call_args_list, [call(180.0)] * 2)
|
||||
summarize_mock.assert_called_once_with()
|
||||
|
||||
def test_six_minute_wait_enters_fill_mode_at_checkpoint(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
counts = {
|
||||
"pending": 5,
|
||||
"recently_downloaded": 0,
|
||||
"unobserved": 0,
|
||||
"unknown": 0,
|
||||
"ready": 0,
|
||||
"total": 5,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(monitor, "summarize", return_value=counts),
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.time.sleep",
|
||||
) as sleep_mock,
|
||||
):
|
||||
entered_fill_mode = monitor.wait_for_dynamic_queue(
|
||||
LARGE_POLYGON_WAIT,
|
||||
"Queue wait complete.",
|
||||
)
|
||||
|
||||
self.assertTrue(entered_fill_mode)
|
||||
self.assertTrue(monitor.filling_open_queue)
|
||||
sleep_mock.assert_called_once_with(180.0)
|
||||
|
||||
def test_fifteen_minute_wait_checks_at_five_and_ten_minutes(self) -> None:
|
||||
monitor = self.make_monitor()
|
||||
counts_with_fourteen_slots = {
|
||||
"pending": 6,
|
||||
"recently_downloaded": 0,
|
||||
"unobserved": 0,
|
||||
"unknown": 0,
|
||||
"ready": 0,
|
||||
"total": 6,
|
||||
}
|
||||
counts_with_fifteen_slots = {
|
||||
**counts_with_fourteen_slots,
|
||||
"pending": 5,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
monitor,
|
||||
"summarize",
|
||||
side_effect=[counts_with_fourteen_slots, counts_with_fifteen_slots],
|
||||
) as summarize_mock,
|
||||
patch(
|
||||
"request_nsrdb_county_polygon_archives.time.sleep",
|
||||
) as sleep_mock,
|
||||
):
|
||||
entered_fill_mode = monitor.wait_for_dynamic_queue(
|
||||
EXCEPTIONAL_POLYGON_WAIT,
|
||||
"Queue wait complete.",
|
||||
)
|
||||
|
||||
self.assertTrue(entered_fill_mode)
|
||||
self.assertTrue(monitor.filling_open_queue)
|
||||
self.assertEqual(sleep_mock.call_args_list, [call(300.0)] * 2)
|
||||
self.assertEqual(summarize_mock.call_count, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user