Skip to content

Commit

Permalink
tests: memoize the generated test known_hosts file
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
evanelias committed Sep 18, 2023
1 parent 7198c0f commit 2442217
Showing 1 changed file with 28 additions and 4 deletions.
32 changes: 28 additions & 4 deletions knownhosts_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 2442217

Please sign in to comment.