215 lines
11 KiB
Python
215 lines
11 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import datetime, timedelta, timezone
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from solar_relay.cli import next_poll_at
|
|
from solar_relay.config import retention_cutoff
|
|
from solar_relay.storage import SCHEMA, Store, StoreError
|
|
from tests.test_status import observed
|
|
|
|
|
|
class ParquetTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.directory.cleanup)
|
|
self.path = Path(self.directory.name) / "archive"
|
|
self.store = Store(self.path)
|
|
self.now = datetime(2026, 11, 1, 12, tzinfo=timezone.utc)
|
|
|
|
def test_real_parquet_has_typed_readings_and_original_payload(self):
|
|
value = observed(self.now)
|
|
value["readings"]["boiler_enabled"] = False
|
|
value["readings"]["boiler_threshold_w"] = 2600
|
|
value["readings"]["boiler_threshold_source"] = "config"
|
|
self.store.add(value)
|
|
path = next(self.path.glob("date=*/samples.parquet"))
|
|
self.assertEqual(path.read_bytes()[:4], b"PAR1")
|
|
row = pq.ParquetFile(path).read().to_pylist()[0]
|
|
self.assertEqual(row["production_ac_w"], 2000)
|
|
self.assertIs(row["boiler_enabled"], False)
|
|
self.assertEqual(row["boiler_threshold_w"], 2600)
|
|
self.assertEqual(row["boiler_threshold_source"], "config")
|
|
self.assertIsNone(row["uploaded_at"])
|
|
self.assertEqual(json.loads(row["payload"]), value)
|
|
self.assertEqual(self.store.status()["format"], "parquet")
|
|
|
|
def test_retention_uses_ack_calendar_days_and_never_deletes_pending(self):
|
|
items = [observed(self.now - timedelta(days=90 + index)) for index in range(4)]
|
|
for item in items:
|
|
self.store.add(item)
|
|
self.store.acknowledge([items[0]["id"]], now=self.now - timedelta(days=31))
|
|
self.store.acknowledge([items[1]["id"]], now=self.now - timedelta(days=30))
|
|
self.store.acknowledge([items[2]["id"]], now=self.now)
|
|
self.assertEqual(self.store.prune(now=self.now), 1)
|
|
self.assertEqual({s["id"] for s in self.store.records()}, {s["id"] for s in items[1:]})
|
|
self.assertEqual(self.store.pending(100), [items[3]])
|
|
self.assertEqual(self.store.prune(now=self.now + timedelta(days=31)), 2)
|
|
self.assertEqual(list(self.store.records()), [items[3]])
|
|
|
|
def test_additional_archive_columns_do_not_change_records_or_acknowledgements(self):
|
|
value = observed(self.now)
|
|
self.store.add(value)
|
|
path = next(self.path.glob("date=*/samples.parquet"))
|
|
table = pq.ParquetFile(path).read()
|
|
table = table.append_column("unused_estimate", pa.array([100.0]))
|
|
pq.write_table(table, path)
|
|
self.assertEqual(self.store.pending(100), [value])
|
|
self.assertEqual(list(self.store.records()), [value])
|
|
self.store.acknowledge([value["id"]], now=self.now)
|
|
self.assertEqual(self.store.pending(100), [])
|
|
self.assertEqual(list(self.store.records()), [value])
|
|
self.assertTrue(pq.ParquetFile(path).schema_arrow.equals(SCHEMA, check_metadata=True))
|
|
|
|
def test_missing_or_mistyped_archive_columns_are_rejected(self):
|
|
value = observed(self.now)
|
|
self.store.add(value)
|
|
path = next(self.path.glob("date=*/samples.parquet"))
|
|
table = pq.ParquetFile(path).read()
|
|
index = table.schema.get_field_index("production_ac_w")
|
|
for invalid in (table.drop_columns(["production_ac_w"]),
|
|
table.set_column(index, "production_ac_w", pa.array(["invalid"]))):
|
|
with self.subTest(schema=invalid.schema):
|
|
pq.write_table(invalid, path)
|
|
with self.assertRaises(StoreError):
|
|
self.store.pending(100)
|
|
|
|
def test_pre_threshold_archive_is_readable_and_upgraded_on_write(self):
|
|
value = observed(self.now)
|
|
self.store.add(value)
|
|
path = next(self.path.glob("date=*/samples.parquet"))
|
|
legacy = pq.ParquetFile(path).read().drop_columns(["boiler_threshold_w", "boiler_threshold_source"])
|
|
pq.write_table(legacy, path)
|
|
self.assertEqual(self.store.pending(100), [value])
|
|
self.assertEqual(list(self.store.records()), [value])
|
|
self.store.acknowledge([value["id"]], now=self.now)
|
|
self.assertTrue(pq.ParquetFile(path).schema_arrow.equals(SCHEMA, check_metadata=True))
|
|
self.assertIsNone(pq.ParquetFile(path).read().to_pylist()[0]["boiler_threshold_w"])
|
|
|
|
def test_retention_boundary_is_local_midnight_regardless_of_execution_time(self):
|
|
zone = ZoneInfo("Europe/Zurich")
|
|
for hour in (0, 12, 23):
|
|
now = datetime(2026, 11, 2, hour, 30, tzinfo=zone)
|
|
self.assertEqual(retention_cutoff(now, 30, zone.key), datetime(2026, 10, 2, 22, tzinfo=timezone.utc))
|
|
with self.assertRaisesRegex(ValueError, "timezone"):
|
|
retention_cutoff(datetime(2026, 11, 2), 30, zone.key)
|
|
|
|
def test_prune_uses_the_configured_timezone_at_the_date_boundary(self):
|
|
store = Store(self.path, timezone_name="America/New_York")
|
|
values = [observed(self.now - timedelta(days=60 + offset)) for offset in range(2)]
|
|
for value in values:
|
|
store.add(value)
|
|
store.acknowledge([values[0]["id"]], now=datetime(2026, 10, 2, 3, 59, tzinfo=timezone.utc))
|
|
store.acknowledge([values[1]["id"]], now=datetime(2026, 10, 2, 4, tzinfo=timezone.utc))
|
|
self.assertEqual(store.prune(now=datetime(2026, 11, 2, 1, 30, tzinfo=timezone.utc)), 1)
|
|
self.assertEqual(list(store.records()), [values[1]])
|
|
|
|
def test_prune_keeps_the_entire_cutoff_acknowledgement_day(self):
|
|
zone = ZoneInfo("Europe/Zurich")
|
|
start = datetime(2026, 10, 3, tzinfo=zone)
|
|
values = [observed(start - timedelta(days=10) + timedelta(minutes=offset)) for offset in (0, 5)]
|
|
for value, ack in zip(values, (start, start + timedelta(hours=23, minutes=55))):
|
|
self.store.add(value)
|
|
self.store.acknowledge([value["id"]], now=ack)
|
|
self.assertEqual(self.store.prune(now=datetime(2026, 11, 2, 23, 59, tzinfo=zone)), 0)
|
|
self.assertEqual(self.store.prune(now=datetime(2026, 11, 3, tzinfo=zone)), 2)
|
|
|
|
def test_prune_removes_whole_local_days_across_utc_partitions(self):
|
|
zone = ZoneInfo("Europe/Zurich")
|
|
first = datetime(2026, 10, 3, tzinfo=zone)
|
|
expired = [observed(first - timedelta(hours=23, minutes=55)), observed(first - timedelta(minutes=5))]
|
|
retained = [observed(first), observed(first + timedelta(hours=23, minutes=55))]
|
|
for value in expired + retained:
|
|
self.store.add(value)
|
|
self.store.acknowledge([value["id"] for value in expired], now=first - timedelta(minutes=1))
|
|
self.store.acknowledge([value["id"] for value in retained], now=first + timedelta(days=1))
|
|
self.assertEqual(self.store.prune(now=datetime(2026, 11, 2, 22, tzinfo=zone)), 2)
|
|
self.assertEqual(list(self.store.records()), retained)
|
|
|
|
def test_one_pending_record_protects_its_entire_local_day_across_partitions(self):
|
|
zone = ZoneInfo("Europe/Zurich")
|
|
day = datetime(2026, 10, 2, tzinfo=zone)
|
|
values = [observed(day + timedelta(minutes=5)), observed(day + timedelta(hours=23, minutes=55))]
|
|
for value in values:
|
|
self.store.add(value)
|
|
self.store.acknowledge([values[0]["id"]], now=day + timedelta(hours=23, minutes=59))
|
|
self.assertEqual(self.store.prune(now=datetime(2026, 11, 2, 22, tzinfo=zone)), 0)
|
|
self.assertEqual(list(self.store.records()), values)
|
|
|
|
def test_daylight_saving_retention_boundaries_preserve_complete_local_days(self):
|
|
zone = ZoneInfo("Europe/Zurich")
|
|
for day, expected_hours in [(datetime(2026, 3, 29, tzinfo=zone), 23),
|
|
(datetime(2026, 10, 25, tzinfo=zone), 25)]:
|
|
with self.subTest(day=day):
|
|
now = day + timedelta(days=30, hours=12)
|
|
cutoff = retention_cutoff(now, 30, zone.key)
|
|
self.assertEqual(cutoff, day.astimezone(timezone.utc))
|
|
next_day = (day + timedelta(days=1)).astimezone(timezone.utc)
|
|
self.assertEqual((next_day - cutoff).total_seconds(), expected_hours * 3600)
|
|
|
|
def test_replay_cancels_expiry_until_the_next_acknowledgement(self):
|
|
value = observed(self.now)
|
|
self.store.add(value)
|
|
self.store.acknowledge([value["id"]], now=self.now - timedelta(days=40))
|
|
self.store.replay()
|
|
self.assertEqual(self.store.prune(now=self.now), 0)
|
|
self.store.acknowledge([value["id"]], now=self.now)
|
|
self.assertEqual(self.store.prune(now=self.now), 0)
|
|
|
|
def test_interrupted_atomic_write_preserves_the_previous_partition(self):
|
|
first, second = observed(self.now), observed(self.now + timedelta(minutes=5))
|
|
self.store.add(first)
|
|
with patch("solar_relay.storage.os.replace", side_effect=OSError("simulated write failure")):
|
|
with self.assertRaises(OSError):
|
|
self.store.add(second)
|
|
self.assertEqual(list(Store(self.path).records()), [first])
|
|
self.assertEqual(list(self.path.glob("date=*/.samples-*.tmp")), [])
|
|
|
|
def test_two_store_instances_cannot_lose_concurrent_records(self):
|
|
other = Store(self.path)
|
|
values = [observed(self.now + timedelta(minutes=index * 5)) for index in range(12)]
|
|
def add(pair):
|
|
index, value = pair
|
|
(other if index % 2 else self.store).add(value)
|
|
with ThreadPoolExecutor(max_workers=2) as workers:
|
|
list(workers.map(add, enumerate(values)))
|
|
self.assertEqual({value["id"] for value in self.store.records()}, {value["id"] for value in values})
|
|
|
|
def test_pruning_an_entire_day_removes_its_parquet_file(self):
|
|
value = observed(self.now - timedelta(days=60))
|
|
self.store.add(value)
|
|
self.store.acknowledge([value["id"]], now=self.now - timedelta(days=31))
|
|
self.assertEqual(self.store.prune(now=self.now), 1)
|
|
self.assertEqual(list(self.path.glob("date=*/samples.parquet")), [])
|
|
self.assertEqual(self.store.status()["samples"], 0)
|
|
|
|
def test_corrupt_archive_fails_without_silently_dropping_records(self):
|
|
value = observed(self.now)
|
|
self.store.add(value)
|
|
next(self.path.glob("date=*/samples.parquet")).write_bytes(b"broken")
|
|
with self.assertRaises(StoreError):
|
|
self.store.pending(100)
|
|
|
|
def test_archive_path_requires_a_directory(self):
|
|
occupied = self.path.parent / "existing-file"
|
|
occupied.write_bytes(b"existing data")
|
|
with self.assertRaises(StoreError):
|
|
Store(occupied)
|
|
self.assertEqual(occupied.read_bytes(), b"existing data")
|
|
|
|
def test_next_poll_is_aligned_to_the_five_minute_clock(self):
|
|
self.assertEqual(next_poll_at(301, 300), 600)
|
|
self.assertEqual(next_poll_at(599.5, 300), 600)
|
|
self.assertEqual(next_poll_at(600, 300), 900)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|