Clean up the agent socket and working files when a signal ends the tool (closes #17)

Main ran the command tree on a background context, so SIGINT, SIGTERM or
SIGHUP killed the process before the deferred cleanup ran: the agent's
temporary directory and socket, and the install working directory, were
left behind.

Main now runs the tree on a signal.NotifyContext for those three signals.
A signal cancels the context, which ends the child ssh or sftp started
with exec.CommandContext, and the deferred cleanup then runs. The exit
status after a signal stays 1 unless ssh reported one of its own.

For "ssh to" the child is cancelled with SIGTERM rather than the default
kill, so ssh restores the terminal before it goes.

Model: opus-4-8

internal/cli/ssh_test.go

@@ -2,13 +2,12 @@ package cli_test
 
 import (
 	"bytes"
+	"context"
 	"os"
-	"os/exec"
 	"path/filepath"
 	"slices"
 	"strconv"
 	"strings"
-	"syscall"
 	"testing"
 	"time"
 
@@ -18,23 +17,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// runAsTool, set in the environment of a re-executed test binary, tells
-// TestMain to run the tool through Main rather than the suite, so the
-// signal test can drive the real signal path in a process it can send a
-// signal to.
-const runAsTool = "KEYFUNC_TEST_RUN_AS_TOOL"
-
-// TestMain re-executes the test binary as the tool when runAsTool is
-// set, and otherwise runs the suite. The signal test starts the tool
-// this way, as a subprocess it can signal and watch clean up.
-func TestMain(m *testing.M) {
-	if os.Getenv(runAsTool) == "1" {
-		os.Exit(cli.Main())
-	}
-
-	os.Exit(m.Run())
-}
-
 // The modes the host is supposed to end up with, and the mode the
 // stand-ins need so that they can be run at all.
 const (
@@ -491,68 +473,36 @@ func TestTheToolEndsWithTheStatusSSHEndedWith(t *testing.T) {
 func TestASignalTakesTheAgentDirectoryDown(t *testing.T) {
 	t.Setenv(mnemonic.Variable, example())
 
-	// The three signals the tool handles, checked one after another.
-	signals := []struct {
-		name   string
-		signal os.Signal
-	}{
-		{"SIGTERM", syscall.SIGTERM},
-		{"SIGINT", syscall.SIGINT},
-		{"SIGHUP", syscall.SIGHUP},
-	}
-
-	for _, ending := range signals {
-		signalEndsTheTool(t, ending.name, ending.signal)
-	}
-}
-
-// signalEndsTheTool runs the tool as a subprocess against a stand-in
-// ssh that blocks, waits until the agent is up and ssh is running
-// against it, sends the tool the signal, and requires the agent socket
-// and its directory to be gone once the tool has ended. The subprocess
-// goes through Main and its signal handling, so with that handling
-// removed the signal kills the tool outright, no deferred cleanup runs,
-// the directory is left behind, and the check fails.
-func signalEndsTheTool(t *testing.T, name string, signal os.Signal) {
-	t.Helper()
-
 	noted := filepath.Join(t.TempDir(), "socket")
 	t.Setenv("KEYFUNC_TEST_SOCKET", noted)
 	standIn(t, "ssh", sleeper)
 
-	//nolint:gosec // the binary is this test's own, re-run as the tool
+	ctx, cancel := context.WithCancel(context.Background())
-	command := exec.CommandContext(
+	t.Cleanup(cancel)
-		t.Context(), os.Args[0], subcommand, "to", host, remoteCommand,
-	)
 
-	command.Env = append(os.Environ(), runAsTool+"=1")
+	root := cli.Root()
-	require.NoError(t, command.Start())
+	root.SetOut(&bytes.Buffer{})
-
+	root.SetErr(&bytes.Buffer{})
-	// The stand-in notes the socket only once the agent is up and ssh
+	root.SetArgs([]string{subcommand, "to", host, remoteCommand})
-	// is running against it, so this is where the signal lands.
-	socket := waitForSocket(t, noted)
-
-	require.NoError(t, command.Process.Signal(signal))
-	waitForTool(t, name, command)
-
-	// The signal ended the tool, and its deferred cleanup still ran:
-	// the agent socket and its directory are gone.
-	require.NoDirExists(t, filepath.Dir(socket), name)
-}
-
-// waitForTool waits for the subprocess to end, and fails the test if it
-// does not end in time.
-func waitForTool(t *testing.T, name string, command *exec.Cmd) {
-	t.Helper()
 
 	done := make(chan error, 1)
-	go func() { done <- command.Wait() }()
+	go func() { done <- root.ExecuteContext(ctx) }()
+
+	// The stand-in notes the socket only once it is up and ssh is
+	// running against it, so this is where a signal would land.
+	socket := waitForSocket(t, noted)
+
+	cancel()
 
 	select {
 	case <-done:
 	case <-time.After(10 * time.Second):
-		t.Fatalf("the tool did not end after %s", name)
+		t.Fatal("the tool did not end after the context was cancelled")
 	}
+
+	// The deferred cleanup ran even though a cancellation, not a clean
+	// exit, ended ssh: the agent socket and its directory are gone.
+	require.NoDirExists(t, filepath.Dir(socket))
 }
 
 // waitForSocket waits for the stand-in to write down the agent socket