From 77d3704f847d3ef32a33a0c18b330fb5e2534cb1 Mon Sep 17 00:00:00 2001
From: Patrick Roy <roypat@amazon.co.uk>
Date: Fri, 13 Dec 2024 11:23:15 +0000
Subject: [PATCH] refactor(test): improve retry look in check_guest_connections

Use `tenancy.Retrying` instead of the home-grown retry loop. This has
the advantage that we don't simply give up after 3 attempts, but instead
raise a meaningful error message (whereas before I guess we would just
fail further down the function when trying to do something that assumes
the file exists).

Signed-off-by: Patrick Roy <roypat@amazon.co.uk>
---
 tests/framework/utils_vsock.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tests/framework/utils_vsock.py b/tests/framework/utils_vsock.py
index 244850c3496..9561c1c26f2 100644
--- a/tests/framework/utils_vsock.py
+++ b/tests/framework/utils_vsock.py
@@ -11,6 +11,8 @@
 from subprocess import Popen
 from threading import Thread
 
+from tenacity import Retrying, stop_after_attempt, wait_fixed
+
 ECHO_SERVER_PORT = 5252
 SERVER_ACCEPT_BACKLOG = 128
 TEST_CONNECTION_COUNT = 50
@@ -143,13 +145,17 @@ def check_guest_connections(vm, server_port_path, blob_path, blob_hash):
     )
 
     try:
+        # Give socat a bit of time to create the socket
+        for attempt in Retrying(
+            wait=wait_fixed(0.2),
+            stop=stop_after_attempt(3),
+            reraise=True,
+        ):
+            with attempt:
+                assert Path(server_port_path).exists()
+
         # Link the listening Unix socket into the VM's jail, so that
         # Firecracker can connect to it.
-        attempt = 0
-        # But 1st, give socat a bit of time to create the socket
-        while not Path(server_port_path).exists() and attempt < 3:
-            time.sleep(0.2)
-            attempt += 1
         vm.create_jailed_resource(server_port_path)
 
         # Increase maximum process count for the ssh service.