refactor migration transition test

Signed-off-by: William Henderson <william.henderson@nutanix.com>

1 files changed, 8 insertions, 2 deletions

diff --git a/test/py/test_migration.py b/test/py/test_migration.py
index 7370b52..b320e76 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -259,9 +259,15 @@ def test_migration_shortest_state_transition_paths():
         while len(queue) > 0:
             (curr, prev) = queue.popleft()
             back[curr] = prev
+
+            # Intermediate states cannot be saving states, so if our current
+            # node is not the start state and it is a saving state, it is only
+            # allowed to be an end state so we don't explore its neighbours.
+            if curr != source and curr in saving_states:
+                continue
+
             for nxt in E[curr]:
-                if back[nxt] is None \
-                        and (curr == source or curr not in saving_states):
+                if back[nxt] is None:
                     queue.append((nxt, curr))
 
         # Iterate over the states