From 2442217a08f2072ac0640054256d10f5b6268978 Mon Sep 17 00:00:00 2001
From: Evan Elias <evan@skeema.net>
Date: Mon, 18 Sep 2023 15:30:56 -0400
Subject: [PATCH] tests: memoize the generated test known_hosts file

Previously, several test functions each wrote their own test known_hosts file,
generating different random keys each time. This is slow and CPU-intensive.

This commit changes the test logic to generate random keys once per overall
test process, and re-use those test known_hosts contents across multiple test
functions.
---
 knownhosts_test.go | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/knownhosts_test.go b/knownhosts_test.go
index f48d2b2..329a81b 100644
--- a/knownhosts_test.go
+++ b/knownhosts_test.go
@@ -16,7 +16,7 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	khPath := writeTestKnownHosts(t)
+	khPath := getTestKnownHosts(t)
 
 	// Valid path should return a callback and no error; callback should be usable
 	// in ssh.ClientConfig.HostKeyCallback
@@ -35,7 +35,7 @@ func TestNew(t *testing.T) {
 }
 
 func TestHostKeyAlgorithms(t *testing.T) {
-	khPath := writeTestKnownHosts(t)
+	khPath := getTestKnownHosts(t)
 	kh, err := New(khPath)
 	if err != nil {
 		t.Fatalf("Unexpected error from New: %v", err)
@@ -67,7 +67,7 @@ func TestHostKeyAlgorithms(t *testing.T) {
 }
 
 func TestIsHostKeyChanged(t *testing.T) {
-	khPath := writeTestKnownHosts(t)
+	khPath := getTestKnownHosts(t)
 	kh, err := New(khPath)
 	if err != nil {
 		t.Fatalf("Unexpected error from New: %v", err)
@@ -104,7 +104,7 @@ func TestIsHostKeyChanged(t *testing.T) {
 }
 
 func TestIsHostUnknown(t *testing.T) {
-	khPath := writeTestKnownHosts(t)
+	khPath := getTestKnownHosts(t)
 	kh, err := New(khPath)
 	if err != nil {
 		t.Fatalf("Unexpected error from New: %v", err)
@@ -222,6 +222,30 @@ func TestWriteKnownHost(t *testing.T) {
 	}
 }
 
+var testKnownHostsContents []byte
+
+// getTestKnownHosts returns a path to a test known_hosts file. The file path
+// will differ between test functions, but the contents are always the same,
+// containing keys generated upon the first invocation. The file is removed
+// upon test completion.
+func getTestKnownHosts(t *testing.T) string {
+	// Re-use previously memoized result
+	if len(testKnownHostsContents) > 0 {
+		dir := t.TempDir()
+		khPath := filepath.Join(dir, "known_hosts")
+		if err := os.WriteFile(khPath, testKnownHostsContents, 0600); err != nil {
+			t.Fatalf("Unable to write to %s: %v", khPath, err)
+		}
+		return khPath
+	}
+
+	khPath := writeTestKnownHosts(t)
+	if contents, err := os.ReadFile(khPath); err == nil {
+		testKnownHostsContents = contents
+	}
+	return khPath
+}
+
 // writeTestKnownHosts generates the test known_hosts file and returns the
 // file path to it. The generated file contains several hosts with a mix of
 // key types; each known host has between 1 and 3 different known host keys.