Updated statuses management, tests (#41)

* Updated statuses management, tests

* storage: query: Generalize event loading logic.

* reformat

* storage: query: Rename load_events to prefetch_event_relations.

Co-authored-by: Giacomo Leidi <goodoldpaul@autistici.org>

mobilizon_bots/event/event.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, asdict
-from typing import Optional
+from enum import IntEnum
+from typing import Optional, Set
 
 import arrow
 import tortoise.timezone
@@ -9,6 +10,13 @@ from mobilizon_bots.models.event import Event
 from mobilizon_bots.models.publication import PublicationStatus, Publication
 
 
+class EventPublicationStatus(IntEnum):
+    WAITING = 1
+    FAILED = 2
+    COMPLETED = 3
+    PARTIAL = 4
+
+
 @dataclass
 class MobilizonEvent:
     """Class representing an event retrieved from Mobilizon."""
@@ -22,15 +30,15 @@ class MobilizonEvent:
     thumbnail_link: Optional[str] = None
     location: Optional[str] = None
     publication_time: Optional[dict[str, arrow.Arrow]] = None
-    publication_status: PublicationStatus = PublicationStatus.WAITING
+    status: EventPublicationStatus = EventPublicationStatus.WAITING
 
     def __post_init__(self):
         assert self.begin_datetime.tzinfo == self.end_datetime.tzinfo
         assert self.begin_datetime < self.end_datetime
         if self.publication_time:
-            assert self.publication_status in [
+            assert self.status in [
-                PublicationStatus.COMPLETED,
+                EventPublicationStatus.COMPLETED,
-                PublicationStatus.PARTIAL,
+                EventPublicationStatus.PARTIAL,
             ]
 
     def _fill_template(self, pattern: Template) -> str:
@@ -52,19 +60,20 @@ class MobilizonEvent:
         )
 
     @staticmethod
-    def compute_status(publications: list[Publication]):
+    def compute_status(publications: list[Publication]) -> EventPublicationStatus:
-        unique_statuses = set(pub.status for pub in publications)
+        unique_statuses: Set[PublicationStatus] = set(
-        assert PublicationStatus.PARTIAL not in unique_statuses
+            pub.status for pub in publications
+        )
 
         if PublicationStatus.FAILED in unique_statuses:
-            return PublicationStatus.FAILED
+            return EventPublicationStatus.FAILED
         elif unique_statuses == {
             PublicationStatus.COMPLETED,
             PublicationStatus.WAITING,
         }:
-            return PublicationStatus.PARTIAL
+            return EventPublicationStatus.PARTIAL
         elif len(unique_statuses) == 1:
-            return unique_statuses.pop()
+            return EventPublicationStatus[unique_statuses.pop().name]
 
         raise ValueError(f"Illegal combination of PublicationStatus: {unique_statuses}")
 
@@ -84,7 +93,6 @@ class MobilizonEvent:
             mobilizon_id=event.mobilizon_id,
             thumbnail_link=event.thumbnail_link,
             location=event.location,
-            # TODO: Discuss publications (both time and status)
             publication_time={
                 pub.publisher.name: arrow.get(
                     tortoise.timezone.localtime(value=pub.timestamp, timezone=tz)
@@ -93,5 +101,5 @@ class MobilizonEvent:
             }
             if publication_status != PublicationStatus.WAITING
             else None,
-            publication_status=publication_status,
+            status=publication_status,
         )

mobilizon_bots/event/event_selection_strategies.py

@@ -110,7 +110,8 @@ STRATEGY_NAME_TO_STRATEGY_CLASS = {"next_event": SelectNextEventStrategy}
 
 
 def select_event_to_publish(
-    published_events: List[MobilizonEvent], unpublished_events: List[MobilizonEvent],
+    published_events: List[MobilizonEvent],
+    unpublished_events: List[MobilizonEvent],
 ):
 
     strategy = STRATEGY_NAME_TO_STRATEGY_CLASS[

mobilizon_bots/mobilizon/events.py

@@ -39,7 +39,7 @@ def parse_event(data):
         thumbnail_link=parse_picture(data),
         location=parse_location(data),
         publication_time=None,
-        publication_status=PublicationStatus.WAITING,
+        status=PublicationStatus.WAITING,
     )
 
 

mobilizon_bots/models/publication.py

@@ -7,8 +7,7 @@ from tortoise.models import Model
 class PublicationStatus(IntEnum):
     WAITING = 1
     FAILED = 2
-    PARTIAL = 3
+    COMPLETED = 3
-    COMPLETED = 4
 
 
 class Publication(Model):

mobilizon_bots/publishers/coordinator.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass, field
+from enum import IntEnum
 from typing import List
 
-from mobilizon_bots.event.event import MobilizonEvent, PublicationStatus
+from mobilizon_bots.event.event import MobilizonEvent
 from mobilizon_bots.publishers import get_active_publishers
 from mobilizon_bots.publishers.abstract import AbstractPublisher
 from mobilizon_bots.publishers.exceptions import PublisherError
@@ -10,9 +11,15 @@ from mobilizon_bots.publishers.telegram import TelegramPublisher
 KEY2CLS = {"telegram": TelegramPublisher}
 
 
+class PublisherStatus(IntEnum):
+    WAITING = 1
+    FAILED = 2
+    COMPLETED = 3
+
+
 @dataclass
 class PublisherReport:
-    status: PublicationStatus
+    status: PublisherStatus
     reason: str
     publisher: AbstractPublisher
 
@@ -23,7 +30,7 @@ class PublisherCoordinatorReport:
 
     @property
     def successful(self):
-        return all(r.status == PublicationStatus.COMPLETED for r in self.reports)
+        return all(r.status == PublisherStatus.COMPLETED for r in self.reports)
 
     def __iter__(self):
         return self.reports.__iter__()
@@ -43,7 +50,11 @@ class PublisherCoordinator:
 
     def _make_successful_report(self):
         return [
-            PublisherReport(status=PublicationStatus.COMPLETED, reason="", publisher=p,)
+            PublisherReport(
+                status=PublisherStatus.COMPLETED,
+                reason="",
+                publisher=p,
+            )
             for p in self.publishers
         ]
 
@@ -55,7 +66,9 @@ class PublisherCoordinator:
             except PublisherError as e:
                 failed_publishers_reports.append(
                     PublisherReport(
-                        status=PublicationStatus.FAILED, reason=repr(e), publisher=p,
+                        status=PublisherStatus.FAILED,
+                        reason=repr(e),
+                        publisher=p,
                     )
                 )
         reports = failed_publishers_reports or self._make_successful_report()
@@ -67,7 +80,7 @@ class PublisherCoordinator:
             if not p.are_credentials_valid():
                 invalid_credentials.append(
                     PublisherReport(
-                        status=PublicationStatus.FAILED,
+                        status=PublisherStatus.FAILED,
                         reason="Invalid credentials",
                         publisher=p,
                     )
@@ -75,7 +88,7 @@ class PublisherCoordinator:
             if not p.is_event_valid():
                 invalid_event.append(
                     PublisherReport(
-                        status=PublicationStatus.FAILED,
+                        status=PublisherStatus.FAILED,
                         reason="Invalid event",
                         publisher=p,
                     )
@@ -83,7 +96,7 @@ class PublisherCoordinator:
             if not p.is_message_valid():
                 invalid_msg.append(
                     PublisherReport(
-                        status=PublicationStatus.FAILED,
+                        status=PublisherStatus.FAILED,
                         reason="Invalid message",
                         publisher=p,
                     )

mobilizon_bots/publishers/telegram.py

@@ -42,7 +42,8 @@ class TelegramPublisher(AbstractPublisher):
             err.append("username")
         if err:
             self._log_error(
-                ", ".join(err) + " is/are missing", raise_error=InvalidCredentials,
+                ", ".join(err) + " is/are missing",
+                raise_error=InvalidCredentials,
             )
 
         res = requests.get(f"https://api.telegram.org/bot{token}/getMe")
@@ -50,7 +51,8 @@ class TelegramPublisher(AbstractPublisher):
 
         if not username == data.get("result", {}).get("username"):
             self._log_error(
-                "Found a different bot than the expected one", raise_error=InvalidBot,
+                "Found a different bot than the expected one",
+                raise_error=InvalidBot,
             )
 
     def validate_event(self) -> None:
@@ -63,7 +65,8 @@ class TelegramPublisher(AbstractPublisher):
             res.raise_for_status()
         except requests.exceptions.HTTPError as e:
             self._log_error(
-                f"Server returned invalid data: {str(e)}", raise_error=InvalidResponse,
+                f"Server returned invalid data: {str(e)}",
+                raise_error=InvalidResponse,
             )
 
         try:
@@ -76,7 +79,8 @@ class TelegramPublisher(AbstractPublisher):
 
         if not data.get("ok"):
             self._log_error(
-                f"Invalid request (response: {data})", raise_error=InvalidResponse,
+                f"Invalid request (response: {data})",
+                raise_error=InvalidResponse,
             )
 
         return data

mobilizon_bots/storage/query.py

@@ -2,6 +2,7 @@ import sys
 
 from typing import Iterable, Optional
 
+from tortoise.queryset import QuerySet
 from tortoise.transactions import atomic
 
 from mobilizon_bots.event.event import MobilizonEvent
@@ -17,22 +18,29 @@ from mobilizon_bots.publishers.coordinator import PublisherCoordinatorReport
 CONNECTION_NAME = "models" if "pytest" in sys.modules else None
 
 
+async def prefetch_event_relations(queryset: QuerySet[Event]) -> list[Event]:
+    return (
+        await queryset.prefetch_related("publications__publisher")
+        .order_by("begin_datetime")
+        .distinct()
+    )
+
+
 async def events_with_status(
     statuses: list[PublicationStatus],
 ) -> Iterable[MobilizonEvent]:
     return map(
         MobilizonEvent.from_model,
-        await Event.filter(publications__status__in=statuses)
+        await prefetch_event_relations(Event.filter(publications__status__in=statuses)),
-        .prefetch_related("publications")
-        .prefetch_related("publications__publisher")
-        .order_by("begin_datetime")
-        .distinct(),
     )
 
 
 async def get_published_events() -> Iterable[MobilizonEvent]:
-    return await events_with_status(
+    return map(
-        [PublicationStatus.COMPLETED, PublicationStatus.PARTIAL]
+        MobilizonEvent.from_model,
+        await prefetch_event_relations(
+            Event.filter(publications__status=PublicationStatus.COMPLETED)
+        ),
     )
 
 

tests/conftest.py

@@ -5,7 +5,7 @@ import arrow
 import pytest
 from tortoise.contrib.test import finalizer, initializer
 
-from mobilizon_bots.event.event import MobilizonEvent
+from mobilizon_bots.event.event import MobilizonEvent, EventPublicationStatus
 from mobilizon_bots.models.event import Event
 from mobilizon_bots.models.notification import Notification, NotificationStatus
 from mobilizon_bots.models.publication import Publication, PublicationStatus
@@ -16,6 +16,14 @@ def generate_publication_status(published):
     return PublicationStatus.COMPLETED if published else PublicationStatus.WAITING
 
 
+def generate_event_status(published):
+    return (
+        EventPublicationStatus.COMPLETED
+        if published
+        else EventPublicationStatus.WAITING
+    )
+
+
 def generate_notification_status(published):
     return NotificationStatus.COMPLETED if published else NotificationStatus.WAITING
 
@@ -38,7 +46,7 @@ def event_generator():
             mobilizon_id=mobilizon_id,
             thumbnail_link="http://some_link.com/123.jpg",
             location="location",
-            publication_status=generate_publication_status(published),
+            status=generate_event_status(published),
             publication_time=publication_time
             or (begin_date.shift(days=-1) if published else None),
         )
@@ -115,7 +123,9 @@ def event_model_generator():
 
 @pytest.fixture()
 def publisher_model_generator():
-    def _publisher_model_generator(idx=1,):
+    def _publisher_model_generator(
+        idx=1,
+    ):
         return Publisher(name=f"publisher_{idx}", account_ref=f"account_ref_{idx}")
 
     return _publisher_model_generator

tests/mobilizon/conftest.py

@@ -15,7 +15,10 @@ def mock_mobilizon_success_answer(mobilizon_answer, mobilizon_url):
     with responses.RequestsMock() as rsps:
 
         rsps.add(
-            responses.POST, mobilizon_url, json=mobilizon_answer, status=200,
+            responses.POST,
+            mobilizon_url,
+            json=mobilizon_answer,
+            status=200,
         )
         yield
 
@@ -26,6 +29,8 @@ def mock_mobilizon_failure_answer(mobilizon_url):
     with responses.RequestsMock() as rsps:
 
         rsps.add(
-            responses.POST, mobilizon_url, status=500,
+            responses.POST,
+            mobilizon_url,
+            status=500,
         )
         yield

tests/models/test_event.py

@@ -7,6 +7,7 @@ import tortoise.timezone
 from mobilizon_bots.event.event import MobilizonEvent
 from mobilizon_bots.models.event import Event
 from mobilizon_bots.models.publication import PublicationStatus
+from mobilizon_bots.event.event import EventPublicationStatus
 
 
 @pytest.mark.asyncio
@@ -148,7 +149,7 @@ async def test_mobilizon_event_from_model(
     assert event.thumbnail_link == "thumblink_1"
     assert event.location == "loc_1"
     assert event.publication_time[publisher_model.name] == publication.timestamp
-    assert event.publication_status == PublicationStatus.PARTIAL
+    assert event.status == EventPublicationStatus.PARTIAL
 
 
 @pytest.mark.asyncio
@@ -209,7 +210,7 @@ async def test_mobilizon_event_compute_status_partial(
 
     assert (
         MobilizonEvent.compute_status([publication, publication_2])
-        == PublicationStatus.PARTIAL
+        == EventPublicationStatus.PARTIAL
     )
 
 
@@ -240,5 +241,5 @@ async def test_mobilizon_event_compute_status_waiting(
 
     assert (
         MobilizonEvent.compute_status([publication, publication_2])
-        == PublicationStatus.WAITING
+        == EventPublicationStatus.WAITING
     )

tests/storage/test_query.py

@@ -81,7 +81,7 @@ async def test_get_published_events(
     published_events = list(await get_published_events())
     assert len(published_events) == 1
 
-    assert published_events[0].name == events[0].name
+    assert published_events[0].mobilizon_id == events[0].mobilizon_id
 
     assert published_events[0].begin_datetime == arrow.get(today)
 
@@ -94,13 +94,13 @@ async def test_get_unpublished_events(
         publisher_model_generator, publication_model_generator, event_model_generator
     )
 
-    published_events = list(await get_unpublished_events())
+    unpublished_events = list(await get_unpublished_events())
-    assert len(published_events) == 2
+    assert len(unpublished_events) == 2
 
-    assert published_events[0].name == events[2].name
+    assert unpublished_events[0].mobilizon_id == events[2].mobilizon_id
-    assert published_events[1].name == events[0].name
+    assert unpublished_events[1].mobilizon_id == events[0].mobilizon_id
-    assert published_events[0].begin_datetime == events[2].begin_datetime
+    assert unpublished_events[0].begin_datetime == events[2].begin_datetime
-    assert published_events[1].begin_datetime == events[0].begin_datetime
+    assert unpublished_events[1].begin_datetime == events[0].begin_datetime
 
 
 @pytest.mark.asyncio