tests: fix missing stage sets on state

Note that this change will also break API users (reminder that the API
is NOT semantically versioned). The stage must be set correctly for op
calls to work, primarily for CLI safety, possibly something to disable
in API mode...

Author: Fizzadar

tests/test_api/test_api_deploys.py

@@ -2,6 +2,7 @@
 from pyinfra.api.connect import connect_all, disconnect_all
 from pyinfra.api.deploy import add_deploy, deploy
 from pyinfra.api.operations import run_ops
+from pyinfra.api.state import StateStage
 from pyinfra.operations import server
 
 from ..paramiko_util import PatchSSHTestCase
@@ -15,6 +16,7 @@ def test_deploy(self):
         anotherhost = inventory.get_host("anotherhost")
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
 
         # Enable printing on this test to catch any exceptions in the formatting
         state.print_output = True
@@ -74,6 +76,7 @@ def test_nested_deploy(self):
         somehost = inventory.get_host("somehost")
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
 
         # Enable printing on this test to catch any exceptions in the formatting
         state.print_output = True

tests/test_api/test_api_operations.py

@@ -18,7 +18,7 @@
 from pyinfra.api.exceptions import PyinfraError
 from pyinfra.api.operation import OperationMeta, add_op
 from pyinfra.api.operations import run_ops
-from pyinfra.api.state import StateOperationMeta
+from pyinfra.api.state import StateOperationMeta, StateStage
 from pyinfra.connectors.util import CommandOutput, OutputLine
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra.operations import files, python, server
@@ -44,6 +44,7 @@ def test_op(self):
         anotherhost = inventory.get_host("anotherhost")
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         state.add_callback_handler(BaseStateCallback())
 
         # Enable printing on this test to catch any exceptions in the formatting
@@ -122,6 +123,7 @@ def test_file_upload_op(self):
         inventory = make_inventory()
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Test normal
@@ -197,6 +199,7 @@ def test_file_download_op(self):
         inventory = make_inventory()
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         with patch("pyinfra.operations.files.os.path.isfile", lambda *args, **kwargs: True):
@@ -236,6 +239,7 @@ def test_file_download_op(self):
     def test_function_call_op(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         is_called = []
@@ -257,6 +261,7 @@ def mocked_function(*args, **kwargs):
     def test_run_once_serial_op(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Add a run once op
@@ -280,6 +285,7 @@ def test_run_once_serial_op(self):
     def test_rsync_op(self):
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, files.rsync, "src", "dest", _sudo=True, _sudo_user="root")
@@ -304,6 +310,7 @@ def test_rsync_op(self):
     def test_rsync_op_with_strict_host_key_checking_disabled(self):
         inventory = make_inventory(hosts=(("somehost", {"ssh_strict_host_key_checking": "no"}),))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, files.rsync, "src", "dest", _sudo=True, _sudo_user="root")
@@ -338,6 +345,7 @@ def test_rsync_op_with_strict_host_key_checking_disabled_and_custom_config_file(
             )
         )
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, files.rsync, "src", "dest", _sudo=True, _sudo_user="root")
@@ -365,6 +373,7 @@ def test_rsync_op_with_sanitized_custom_config_file(self):
             hosts=(("somehost", {"ssh_config_file": "/home/me/ssh_test_config && echo hi"}),)
         )
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, files.rsync, "src", "dest", _sudo=True, _sudo_user="root")
@@ -389,6 +398,7 @@ def test_rsync_op_with_sanitized_custom_config_file(self):
     def test_rsync_op_failure(self):
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         with patch("pyinfra.connectors.ssh.which", lambda x: None):
@@ -401,6 +411,7 @@ def test_op_cannot_change_execution_kwargs(self):
         inventory = make_inventory()
 
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
 
         class NoSetDefaultDict(defaultdict):
             def setdefault(self, key, _):
@@ -422,6 +433,7 @@ class TestNestedOperationsApi(PatchSSHTestCase):
     def test_nested_op_api(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
 
         connect_all(state)
 
@@ -458,6 +470,7 @@ class TestOperationFailures(PatchSSHTestCase):
     def test_full_op_fail(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, server.shell, 'echo "hi"')
@@ -484,6 +497,7 @@ def test_full_op_fail(self):
     def test_ignore_errors_op_fail(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         add_op(state, server.shell, 'echo "hi"', _ignore_errors=True)
@@ -514,6 +528,7 @@ class TestOperationOrdering(PatchSSHTestCase):
     def test_cli_op_line_numbers(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         state.current_deploy_filename = __file__
@@ -560,6 +575,7 @@ def test_cli_op_line_numbers(self):
     def test_api_op_line_numbers(self):
         inventory = make_inventory()
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         another_host = inventory.get_host("anotherhost")
@@ -590,6 +606,7 @@ def test_basic_retry_behavior(self, fake_run_command):
         # Create inventory with just one host to simplify testing
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Add operation with retry settings
@@ -649,6 +666,7 @@ def test_retry_max_attempts_failure(self, fake_run_command):
         """
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Add operation with retry settings
@@ -698,6 +716,7 @@ def test_retry_until_condition(self, fake_sleep, fake_run_command):
         # Setup inventory and state using the utility function
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Create a counter to track retry_until calls
@@ -757,6 +776,7 @@ def test_retry_delay(self, fake_sleep, fake_run_command):
         """
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         retry_delay = 5
@@ -799,6 +819,7 @@ def test_retry_until_with_error_handling(self, fake_sleep, fake_run_command):
         """
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Create a retry_until function that raises an exception
@@ -842,6 +863,7 @@ def test_retry_until_with_complex_output_parsing(self, fake_sleep, fake_run_comm
         """
         inventory = make_inventory(hosts=("somehost",))
         state = State(inventory, Config())
+        state.current_stage = StateStage.Prepare
         connect_all(state)
 
         # Track what output we've seen