Fix for server.reboot operation: remove any askpass file at execute time.

Author: Fizzadar

pyinfra/operations/server.py

@@ -68,9 +68,13 @@ def reboot(delay=10, interval=1, reboot_timeout=300, state=None, host=None):
         )
     '''
 
-    # Remove this now, before we reboot the server - if the reboot fails (expected or not)
-    # we'll error if we don't clean this up now. Will simply be re-uploaded if needed later.
-    remove_any_sudo_askpass_file(host)
+    # Remove this now, before we reboot the server - if the reboot fails (expected or
+    # not) we'll error if we don't clean this up now. Will simply be re-uploaded if
+    # needed later.
+    def remove_any_askpass_file(state, host):
+        remove_any_sudo_askpass_file(host)
+
+    yield FunctionCommand(remove_any_askpass_file, (), {})
 
     yield StringCommand('reboot', success_exit_codes=[0, -1])  # -1 being error/disconnected
 

tests/operations/server.reboot/reboot.json

@@ -1,6 +1,7 @@
 {
     "args": [],
     "commands": [
+        ["remove_any_askpass_file", [], {}],
         {
             "command": "reboot",
             "success_exit_codes": [0, -1]