fix async sudo askpass cleanup during disconnect

Author: JakkuSakura

src/pyinfra/api/host.py

@@ -462,12 +462,12 @@ def disconnect(self) -> None:
     async def disconnect_async(self) -> None:
         self._check_state()
 
+        await remove_any_sudo_askpass_file_async(self)
+
         disconnect_func = getattr(self.connector, "disconnect", None)
         if disconnect_func:
             await disconnect_func()
 
-        await remove_any_sudo_askpass_file_async(self)
-
         self.state.trigger_callbacks("host_disconnect", self)
         self.connected = False
         if hasattr(self.state, "active_hosts"):

tests/test_api/test_api_host.py

@@ -1,7 +1,14 @@
 from unittest import TestCase
 
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from pyinfra.api import Config, State
 from pyinfra.api.host import HostData
 
+from ..util import make_inventory
+
 
 class TestHostData(TestCase):
     def test_host_data(self):
@@ -34,3 +41,25 @@ def test_host_data_missing(self):
 
         assert context.exception.args[0] == "Host `somehost` has no data `not-a-key`"
         assert data.get("not-a-key") is None
+
+
+class TestHostDisconnect(TestCase):
+    def test_disconnect_removes_askpass_before_connector_close(self):
+        state = State(make_inventory(), Config())
+        host = state.inventory.get_host("somehost")
+        host.connector_data["sudo_askpass_path"] = "/tmp/askpass"
+
+        calls: list[str] = []
+
+        async def fake_remove(_host):
+            calls.append("remove")
+
+        async def fake_disconnect():
+            calls.append("disconnect")
+
+        host.connector = SimpleNamespace(disconnect=fake_disconnect)
+
+        with patch("pyinfra.api.host.remove_any_sudo_askpass_file_async", fake_remove):
+            asyncio.run(host.disconnect_async())
+
+        assert calls == ["remove", "disconnect"]