From b455d9037f045fe3607999ace15f0a78b2094f5c Mon Sep 17 00:00:00 2001
From: Bryce Lampe <bryce@pulumi.com>
Date: Wed, 27 Mar 2024 15:16:08 -0700
Subject: [PATCH] Add an identity to the agent

---
 examples/main_test.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/examples/main_test.go b/examples/main_test.go
index 981e7ae..4835249 100644
--- a/examples/main_test.go
+++ b/examples/main_test.go
@@ -4,6 +4,8 @@
 package examples
 
 import (
+	"crypto/rand"
+	"crypto/rsa"
 	"net"
 	"os"
 	"path/filepath"
@@ -13,18 +15,18 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	if os.Getenv("SSH_AUTH_SOCK") == "" {
-		sock := sshagent()
-		os.Setenv("SSH_AUTH_SOCK", sock)
-	}
+	sock := sshagent()
+	os.Setenv("SSH_AUTH_SOCK", sock)
 
 	os.Exit(m.Run())
 }
 
+// sshagent crates an in-memory SSH agent with one identity.
 func sshagent() string {
 	dir := os.TempDir()
 	sock := filepath.Join(dir, "test.sock")
 
+	// In case it already exists.
 	_ = os.Remove(sock)
 
 	l, err := net.Listen("unix", sock)
@@ -32,12 +34,18 @@ func sshagent() string {
 		panic(err)
 	}
 
+	a := agent.NewKeyring()
+	key, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		panic(err)
+	}
+	_ = a.Add(agent.AddedKey{PrivateKey: key})
+
 	go func() {
 		conn, err := l.Accept()
 		if err != nil {
 			panic(err)
 		}
-		a := agent.NewKeyring()
 		agent.ServeAgent(a, conn)
 	}()