From fb72c3bee9e8f8319063aa0e3f94052a70b365db Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Sat, 7 Mar 2026 00:09:52 +0100
Subject: [PATCH] Re-try SSH connections in integration tests

Because a VM might be still booting.
---
 internal/tests/ip_endpoint_test.go | 26 ++++++++++++++++++--------
 internal/tests/sshserver_test.go   | 13 ++++++++++++-
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/internal/tests/ip_endpoint_test.go b/internal/tests/ip_endpoint_test.go
index ef6248b..39db437 100644
--- a/internal/tests/ip_endpoint_test.go
+++ b/internal/tests/ip_endpoint_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/avast/retry-go/v5"
 	"github.com/cirruslabs/orchard/internal/tests/devcontroller"
 	"github.com/cirruslabs/orchard/internal/tests/platformdependent"
 	"github.com/cirruslabs/orchard/internal/tests/wait"
@@ -38,19 +39,28 @@ func TestIPEndpoint(t *testing.T) {
 	require.NoError(t, err)
 
 	// Connect to the VM over SSH to make sure the provided IP is valid
-	sshClient, err := ssh.Dial("tcp", ip+":22", &ssh.ClientConfig{
-		User: "admin",
-		Auth: []ssh.AuthMethod{
-			ssh.Password("admin"),
-		},
-		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
-			return nil
-		},
+	sshClient, err := retry.NewWithData[*ssh.Client](
+		retry.Context(t.Context()),
+		retry.DelayType(retry.FixedDelay),
+		retry.Delay(time.Second),
+		retry.Attempts(0),
+	).Do(func() (*ssh.Client, error) {
+		return ssh.Dial("tcp", ip+":22", &ssh.ClientConfig{
+			User: "admin",
+			Auth: []ssh.AuthMethod{
+				ssh.Password("admin"),
+			},
+			HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+				return nil
+			},
+		})
 	})
 	require.NoError(t, err)
+	defer sshClient.Close()
 
 	sshSession, err := sshClient.NewSession()
 	require.NoError(t, err)
+	defer sshSession.Close()
 
 	output, err := sshSession.CombinedOutput("uname -a")
 	require.NoError(t, err)
diff --git a/internal/tests/sshserver_test.go b/internal/tests/sshserver_test.go
index 2becffc..1300018 100644
--- a/internal/tests/sshserver_test.go
+++ b/internal/tests/sshserver_test.go
@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/avast/retry-go/v5"
 	"github.com/cirruslabs/orchard/internal/controller"
 	"github.com/cirruslabs/orchard/internal/tests/devcontroller"
 	"github.com/cirruslabs/orchard/internal/tests/platformdependent"
@@ -82,8 +83,16 @@ func TestSSHServer(t *testing.T) {
 		},
 	})
 	require.NoError(t, err)
+	defer sshClientController.Close()
 
-	netConnVM, err := sshClientController.Dial("tcp", "test-vm:22")
+	netConnVM, err := retry.NewWithData[net.Conn](
+		retry.Context(t.Context()),
+		retry.DelayType(retry.FixedDelay),
+		retry.Delay(time.Second),
+		retry.Attempts(0),
+	).Do(func() (net.Conn, error) {
+		return sshClientController.Dial("tcp", "test-vm:22")
+	})
 	require.NoError(t, err)
 
 	sshConnVM, sshChansVM, sshReqsVM, err := ssh.NewClientConn(netConnVM, "test-vm:22", &ssh.ClientConfig{
@@ -98,9 +107,11 @@ func TestSSHServer(t *testing.T) {
 	require.NoError(t, err)
 
 	sshClientVM := ssh.NewClient(sshConnVM, sshChansVM, sshReqsVM)
+	defer sshClientVM.Close()
 
 	sshSessVM, err := sshClientVM.NewSession()
 	require.NoError(t, err)
+	defer sshSessVM.Close()
 
 	unameBytes, err := sshSessVM.Output("uname -a")
 	require.NoError(t, err)