From 20e63023bb0d3c7e4cb36b91b73e79f51889cc90 Mon Sep 17 00:00:00 2001
From: Andrew Godwin
Date: Tue, 15 Nov 2022 18:30:30 -0700
Subject: Get outbound likes/boosts and their undos working

---
 stator/graph.py  |  4 ++++
 stator/models.py | 14 +++++++++-----
 2 files changed, 13 insertions(+), 5 deletions(-)

(limited to 'stator')

diff --git a/stator/graph.py b/stator/graph.py
index 00ef1c4..ef81d75 100644
--- a/stator/graph.py
+++ b/stator/graph.py
@@ -11,6 +11,7 @@ class StateGraph:
     choices: ClassVar[List[Tuple[object, str]]]
     initial_state: ClassVar["State"]
     terminal_states: ClassVar[Set["State"]]
+    automatic_states: ClassVar[Set["State"]]
 
     def __init_subclass__(cls) -> None:
         # Collect state memebers
@@ -30,6 +31,7 @@ class StateGraph:
                 )
         # Check the graph layout
         terminal_states = set()
+        automatic_states = set()
         initial_state = None
         for state in cls.states.values():
             # Check for multiple initial states
@@ -65,10 +67,12 @@ class StateGraph:
                         raise ValueError(
                             f"State '{state}' does not have a handler method ({state.handler_name})"
                         )
+                    automatic_states.add(state)
         if initial_state is None:
             raise ValueError("The graph has no initial state")
         cls.initial_state = initial_state
         cls.terminal_states = terminal_states
+        cls.automatic_states = automatic_states
         # Generate choices
         cls.choices = [(name, name) for name in cls.states.keys()]
 
diff --git a/stator/models.py b/stator/models.py
index b2cc681..df385dd 100644
--- a/stator/models.py
+++ b/stator/models.py
@@ -105,9 +105,11 @@ class StatorModel(models.Model):
         """
         with transaction.atomic():
             selected = list(
-                cls.objects.filter(state_locked_until__isnull=True, state_ready=True)[
-                    :number
-                ].select_for_update()
+                cls.objects.filter(
+                    state_locked_until__isnull=True,
+                    state_ready=True,
+                    state__in=cls.state_graph.automatic_states,
+                )[:number].select_for_update()
             )
             cls.objects.filter(pk__in=[i.pk for i in selected]).update(
                 state_locked_until=lock_expiry
@@ -144,7 +146,9 @@ class StatorModel(models.Model):
         # If it's a manual progression state don't even try
         # We shouldn't really be here in this case, but it could be a race condition
         if current_state.externally_progressed:
-            print("Externally progressed state!")
+            print(
+                f"Warning: trying to progress externally progressed state {self.state}!"
+            )
             return None
         try:
             next_state = await current_state.handler(self)
@@ -183,7 +187,7 @@ class StatorModel(models.Model):
             state_changed=timezone.now(),
             state_attempted=None,
             state_locked_until=None,
-            state_ready=False,
+            state_ready=True,
         )
 
     atransition_perform = sync_to_async(transition_perform)
-- 
cgit v1.2.3