diff --git a/cmd/pulse-sensor-proxy/main_test.go b/cmd/pulse-sensor-proxy/main_test.go
index 1a7cf84e0e..f0cf44f9d0 100644
--- a/cmd/pulse-sensor-proxy/main_test.go
+++ b/cmd/pulse-sensor-proxy/main_test.go
@@ -319,6 +319,12 @@ func TestHelperProcess(t *testing.T) {
 
 	output := os.Getenv("GO_HELPER_OUTPUT")
 	fmt.Fprint(os.Stdout, output)
+
+	if codeStr := os.Getenv("GO_HELPER_EXIT_CODE"); codeStr != "" {
+		if code, err := strconv.Atoi(codeStr); err == nil {
+			os.Exit(code)
+		}
+	}
 	os.Exit(0)
 }
 
diff --git a/cmd/pulse-sensor-proxy/rpc_test.go b/cmd/pulse-sensor-proxy/rpc_test.go
new file mode 100644
index 0000000000..705763bf60
--- /dev/null
+++ b/cmd/pulse-sensor-proxy/rpc_test.go
@@ -0,0 +1,336 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/rs/zerolog"
+)
+
+type mockKnownHosts struct{}
+
+func (m *mockKnownHosts) Ensure(ctx context.Context, host string) error                   { return nil }
+func (m *mockKnownHosts) EnsureWithPort(ctx context.Context, host string, port int) error { return nil }
+func (m *mockKnownHosts) EnsureWithEntries(ctx context.Context, host string, port int, entries [][]byte) error {
+	return nil
+}
+func (m *mockKnownHosts) Path() string { return "/tmp/known_hosts" }
+
+func TestHandleGetStatusV2(t *testing.T) {
+	tmpDir := t.TempDir()
+	pubKeyPath := filepath.Join(tmpDir, "id_ed25519.pub")
+	os.WriteFile(pubKeyPath, []byte("ssh-ed25519 AAA..."), 0644)
+
+	proxy := &Proxy{
+		sshKeyPath: tmpDir,
+	}
+
+	req := &RPCRequest{}
+	logger := zerolog.New(os.Stderr)
+
+	resp, err := proxy.handleGetStatusV2(context.Background(), req, logger)
+	if err != nil {
+		t.Fatalf("handleGetStatusV2 failed: %v", err)
+	}
+
+	m, ok := resp.(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected map response, got %T", resp)
+	}
+
+	if m["version"] != Version {
+		t.Errorf("expected version %s, got %v", Version, m["version"])
+	}
+	if m["public_key"] != "ssh-ed25519 AAA..." {
+		t.Errorf("unexpected public key: %v", m["public_key"])
+	}
+}
+
+func TestHandleGetTemperatureV2(t *testing.T) {
+	// Mock execLookPath and execCommandFunc
+	origLookPath := execLookPath
+	origExec := execCommandFunc
+	origExecCtx := execCommandContextFunc
+	defer func() {
+		execLookPath = origLookPath
+		execCommandFunc = origExec
+		execCommandContextFunc = origExecCtx
+	}()
+
+	execLookPath = func(file string) (string, error) {
+		return "/usr/bin/" + file, nil
+	}
+
+	mockCmd := func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		args := strings.Join(arg, " ")
+		if name == "sh" && strings.Contains(args, "ssh") {
+			// Mock SSH output with sensor JSON
+			output := `{"cpu": {"temp1": {"temp1_input": 45.0}}}`
+			return mockExecCommand(output)
+		}
+		// For local node
+		if name == "sensors" {
+			output := `{"cpu": {"temp1": {"temp1_input": 45.0}}}`
+			return mockExecCommand(output)
+		}
+		return mockExecCommand("")
+	}
+
+	execCommandFunc = func(name string, arg ...string) *exec.Cmd {
+		return mockCmd(context.Background(), name, arg...)
+	}
+	execCommandContextFunc = func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		return mockCmd(ctx, name, arg...)
+	}
+
+	proxy := &Proxy{
+		config: &Config{
+			AllowedNodes: []string{"node1"},
+		},
+		metrics:    NewProxyMetrics("test"),
+		nodeGate:   newNodeGate(),
+		knownHosts: &mockKnownHosts{},
+	}
+
+	logger := zerolog.New(os.Stderr)
+
+	// Success case
+	req := &RPCRequest{
+		Params: map[string]interface{}{
+			"node": "node1",
+		},
+	}
+
+	resp, err := proxy.handleGetTemperatureV2(context.Background(), req, logger)
+	if err != nil {
+		t.Fatalf("handleGetTemperatureV2 failed: %v", err)
+	}
+
+	m, ok := resp.(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected map response, got %T", resp)
+	}
+	respStr, ok := m["temperature"].(string)
+	if !ok {
+		t.Errorf("expected string temperature, got %T", m["temperature"])
+	}
+	if !strings.Contains(respStr, "45.0") {
+		t.Errorf("expected 45.0 in response, got %s", respStr)
+	}
+
+	// Missing parameter
+	reqMissing := &RPCRequest{Params: map[string]interface{}{}}
+	_, err = proxy.handleGetTemperatureV2(context.Background(), reqMissing, logger)
+	if err == nil {
+		t.Error("expected error for missing node param")
+	}
+
+	// Invalid node (bad format)
+	reqInvalid := &RPCRequest{
+		Params: map[string]interface{}{
+			"node": "invalid node!", // Space should make it invalid
+		},
+	}
+	_, err = proxy.handleGetTemperatureV2(context.Background(), reqInvalid, logger)
+	if err == nil {
+		t.Error("expected error for invalid node")
+	}
+}
+
+func TestHandleRegisterNodesV2(t *testing.T) {
+	// Mock pvecm
+	origLookPath := execLookPath
+	origExec := execCommandFunc
+	origExecCtx := execCommandContextFunc
+	defer func() {
+		execLookPath = origLookPath
+		execCommandFunc = origExec
+		execCommandContextFunc = origExecCtx
+	}()
+
+	execLookPath = func(file string) (string, error) {
+		if file == "pvecm" {
+			return "/usr/sbin/pvecm", nil
+		}
+		return "", fmt.Errorf("not found")
+	}
+
+	mockCmd := func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		args := strings.Join(arg, " ")
+		if name == "pvecm" && args == "status" {
+			// Mock pvecm status output
+			output := "\nCluster information\n-------------------\nName: cluster\nConfig Version: 1\nTransport: knet\nSecure auth: on\n\nQuorum information\n------------------\nDate: Sun Jan 1 00:00:00 2023\nQuorum provider: corosync_votequorum\nNodes: 1\nNode ID: 0x00000001\nRing ID: 1.1\nQuorate: Yes\n\nVotequorum information\n----------------------\nExpected votes: 1\nHighest expected: 1\nTotal votes: 1\nQuorum: 1\nFlags: Quorate \n\nMembership information\n----------------------\n    Nodeid      Votes Name\n0x00000001          1 10.0.0.1 (local)\n"
+			return mockExecCommand(output)
+		}
+
+		// Mock SSH connection test
+		if name == "sh" && strings.Contains(args, "ssh") {
+			// Success
+			return mockExecCommand("")
+		}
+		return mockExecCommand("")
+	}
+
+	execCommandFunc = func(name string, arg ...string) *exec.Cmd {
+		return mockCmd(context.Background(), name, arg...)
+	}
+	execCommandContextFunc = func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		return mockCmd(ctx, name, arg...)
+	}
+
+	tmpDir := t.TempDir()
+	proxy := &Proxy{
+		sshKeyPath: tmpDir,
+		metrics:    NewProxyMetrics("test"),
+		knownHosts: &mockKnownHosts{},
+		config:     &Config{},
+	}
+
+	logger := zerolog.New(os.Stderr)
+	req := &RPCRequest{}
+
+	resp, err := proxy.handleRegisterNodesV2(context.Background(), req, logger)
+	if err != nil {
+		t.Fatalf("handleRegisterNodesV2 failed: %v", err)
+	}
+
+	m, ok := resp.(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected map response, got %T", resp)
+	}
+
+	nodes, ok := m["nodes"].([]map[string]interface{})
+	if !ok {
+		t.Fatalf("expected nodes list, got %T", m["nodes"])
+	}
+
+	if len(nodes) != 1 {
+		t.Errorf("expected 1 node, got %d", len(nodes))
+	}
+	if nodes[0]["name"] != "10.0.0.1" {
+		t.Errorf("expected node name 10.0.0.1, got %v", nodes[0]["name"])
+	}
+	if nodes[0]["ssh_ready"] != true {
+		t.Errorf("expected ssh_ready true, got %v, error: %v", nodes[0]["ssh_ready"], nodes[0]["error"])
+	}
+}
+
+func TestHandleConnection_InvalidJSON(t *testing.T) {
+	// Mock extractPeerCredentials
+	origExtract := extractPeerCredentials
+	defer func() { extractPeerCredentials = origExtract }()
+	extractPeerCredentials = func(conn net.Conn) (*peerCredentials, error) {
+		return &peerCredentials{uid: 1000, gid: 1000, pid: 123}, nil
+	}
+
+	metrics := NewProxyMetrics("test")
+	proxy := &Proxy{
+		metrics: metrics,
+		config: &Config{
+			AllowedPeerUIDs: []uint32{1000},
+			RateLimit: &RateLimitConfig{
+				PerPeerIntervalMs: 1,
+				PerPeerBurst:      10,
+			},
+		},
+		allowedPeerUIDs: map[uint32]struct{}{1000: {}},
+		peerCapabilities: map[uint32]Capability{
+			1000: 3, // LegacyAll
+		},
+		rateLimiter: newRateLimiter(metrics, &RateLimitConfig{
+			PerPeerIntervalMs: 1,
+			PerPeerBurst:      10,
+		}, nil, nil),
+	}
+
+	client, server := net.Pipe()
+
+	go func() {
+		defer client.Close()
+		client.Write([]byte("invalid json\n"))
+	}()
+
+	// handleConnection handles the connection and returns
+	proxy.handleConnection(server)
+	server.Close()
+}
+
+func TestHandleEnsureClusterKeysV2(t *testing.T) {
+	// Mock pvecm and ssh
+	origLookPath := execLookPath
+	origExec := execCommandFunc
+	defer func() {
+		execLookPath = origLookPath
+		execCommandFunc = origExec
+	}()
+
+	execLookPath = func(file string) (string, error) {
+		return "/usr/bin/" + file, nil
+	}
+
+	execCommandFunc = func(name string, arg ...string) *exec.Cmd {
+		args := strings.Join(arg, " ")
+		if name == "pvecm" && args == "status" {
+			output := "\nCluster information\n-------------------\nName: cluster\nConfig Version: 1\nTransport: knet\nSecure auth: on\n\nQuorum information\n------------------\nDate: Sun Jan 1 00:00:00 2023\nQuorum provider: corosync_votequorum\nNodes: 1\nNode ID: 0x00000001\nRing ID: 1.1\nQuorate: Yes\n\nVotequorum information\n----------------------\nExpected votes: 1\nHighest expected: 1\nTotal votes: 1\nQuorum: 1\nFlags: Quorate \n\nMembership information\n----------------------\n    Nodeid      Votes Name\n0x00000001          1 10.0.0.1 (local)\n"
+			return mockExecCommand(output)
+		}
+
+		if name == "sh" && strings.Contains(args, "ssh") {
+			// Check if key exists (grep) - return empty (not found) to trigger adding
+			if strings.Contains(args, "grep -F") {
+				return mockExecCommandWithExitCode("", 1) // 1 = not found
+			}
+			// ensureTempWrapper or key push
+			return mockExecCommand("")
+		}
+		return mockExecCommand("")
+	}
+
+	tmpDir := t.TempDir()
+	// Create SSH keys in tmpDir (for push)
+	os.WriteFile(filepath.Join(tmpDir, "id_ed25519.pub"), []byte("ssh-ed25519 KEY"), 0644)
+
+	proxy := &Proxy{
+		sshKeyPath: tmpDir,
+		metrics:    NewProxyMetrics("test"),
+		knownHosts: &mockKnownHosts{},
+		config: &Config{
+			AllowedSourceSubnets: []string{"10.0.0.0/24"},
+		},
+	}
+
+	logger := zerolog.New(os.Stderr)
+	req := &RPCRequest{}
+
+	resp, err := proxy.handleEnsureClusterKeysV2(context.Background(), req, logger)
+	if err != nil {
+		t.Fatalf("handleEnsureClusterKeysV2 failed: %v", err)
+	}
+
+	m, ok := resp.(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected map response, got %T", resp)
+	}
+
+	if m["success_count"] != 1 {
+		t.Errorf("expected 1 success, got %v", m["success_count"])
+	}
+}
+
+// mockExecCommandWithExitCode helper needed for grep fail
+func mockExecCommandWithExitCode(output string, exitCode int) *exec.Cmd {
+	cs := []string{"-test.run=TestHelperProcess", "--", output}
+	cmd := exec.Command(os.Args[0], cs...)
+	cmd.Env = []string{
+		"GO_WANT_HELPER_PROCESS=1",
+		"GO_HELPER_OUTPUT=" + output,
+		fmt.Sprintf("GO_HELPER_EXIT_CODE=%d", exitCode),
+	}
+	return cmd
+}
diff --git a/cmd/pulse-sensor-proxy/ssh.go b/cmd/pulse-sensor-proxy/ssh.go
index f971d63df0..bd8043d110 100644
--- a/cmd/pulse-sensor-proxy/ssh.go
+++ b/cmd/pulse-sensor-proxy/ssh.go
@@ -27,7 +27,11 @@ var netInterfaces = net.Interfaces
 var execLookPath = exec.LookPath
 
 // Variable for testing to mock os.Hostname
-var osHostname = os.Hostname
+// Variables for testing
+var (
+	osHostname = os.Hostname
+	osStat     = os.Stat
+)
 
 // Variable for testing to mock exec.Command (for simple output)
 var execCommandFunc = exec.Command
@@ -382,7 +386,7 @@ func (p *Proxy) ensureHostKey(node string) error {
 		if err := p.ensureHostKeyFromProxmox(ctx, node); err == nil {
 			return nil
 		} else {
-			if p.config.RequireProxmoxHostkeys {
+			if p.config != nil && p.config.RequireProxmoxHostkeys {
 				return err
 			}
 			log.Warn().
@@ -390,7 +394,7 @@ func (p *Proxy) ensureHostKey(node string) error {
 				Err(err).
 				Msg("Failed to load host key from Proxmox; falling back to ssh-keyscan")
 		}
-	} else if p.config.RequireProxmoxHostkeys {
+	} else if p.config != nil && p.config.RequireProxmoxHostkeys {
 		return fmt.Errorf("require_proxmox_hostkeys enabled but not running on Proxmox host")
 	}
 
@@ -969,7 +973,8 @@ func isProxmoxHost() bool {
 		return true
 	}
 	// Check for /etc/pve directory
-	if info, err := os.Stat("/etc/pve"); err == nil && info.IsDir() {
+	// Check for /etc/pve directory
+	if info, err := osStat("/etc/pve"); err == nil && info.IsDir() {
 		return true
 	}
 	return false
diff --git a/cmd/pulse-sensor-proxy/ssh_test.go b/cmd/pulse-sensor-proxy/ssh_test.go
index 3b76fa2518..4cc9782026 100644
--- a/cmd/pulse-sensor-proxy/ssh_test.go
+++ b/cmd/pulse-sensor-proxy/ssh_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net"
 	"os"
@@ -11,6 +12,8 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+
+	"github.com/rcourtman/pulse-go-rewrite/internal/ssh/knownhosts"
 )
 
 func setupTempWrapper(t *testing.T) (scriptPath, thermalFile, binDir, baseDir string) {
@@ -707,3 +710,308 @@ func TestDiscoverLocalHostAddresses_InterfaceEdgeCases(t *testing.T) {
 	// But we can check that it didn't crash.
 	t.Logf("Got addresses: %v", addrs)
 }
+
+func TestLoadProxmoxHostKeys(t *testing.T) {
+	tmpDir := t.TempDir()
+	hostsFile := filepath.Join(tmpDir, "known_hosts")
+
+	// Override path
+	origPath := proxmoxClusterKnownHostsPath
+	defer func() { proxmoxClusterKnownHostsPath = origPath }()
+	proxmoxClusterKnownHostsPath = hostsFile
+
+	// Create dummy file
+	content := `
+# Comment
+node1 ssh-ed25519 AAA...
+node2 ssh-rsa BBB... comment
+invalid line
+node1 ssh-rsa CCC...
+[node3]:2222 ssh-ed25519 DDD...
+`
+	if err := os.WriteFile(hostsFile, []byte(content), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	// Test success
+	keys, err := loadProxmoxHostKeys("node1")
+	if err != nil {
+		t.Fatalf("loadProxmoxHostKeys failed: %v", err)
+	}
+
+	if len(keys) != 2 {
+		t.Errorf("expected 2 keys for node1, got %d", len(keys))
+	}
+
+	if !strings.Contains(string(keys[0]), "node1 ssh-ed25519 AAA...") {
+		t.Errorf("unexpected key content: %s", keys[0])
+	}
+
+	// Test node2 with comment
+	keys2, err := loadProxmoxHostKeys("node2")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(keys2) != 1 {
+		t.Errorf("expected 1 key for node2")
+	}
+	// Normalization puts comment at end
+	if !strings.Contains(string(keys2[0]), "comment") {
+		t.Errorf("expected comment in key: %s", keys2[0])
+	}
+
+	// Test not found
+	_, err = loadProxmoxHostKeys("unknown")
+	if err == nil {
+		t.Error("expected error for unknown host")
+	}
+
+	// Test file missing
+	proxmoxClusterKnownHostsPath = filepath.Join(tmpDir, "missing")
+	_, err = loadProxmoxHostKeys("node1")
+	if err == nil {
+		t.Error("expected error for missing file")
+	}
+}
+
+func TestBuildAuthorizedKey(t *testing.T) {
+	proxy := &Proxy{
+		config: &Config{
+			AllowedSourceSubnets: []string{"10.0.0.1/32", "192.168.1.0/24"},
+		},
+	}
+
+	pubKey := "ssh-ed25519 AAA..."
+	line, err := proxy.buildAuthorizedKey(pubKey)
+	if err != nil {
+		t.Fatalf("buildAuthorizedKey failed: %v", err)
+	}
+
+	expectedStart := `from="10.0.0.1/32,192.168.1.0/24",command="/usr/local/libexec/pulse-sensor-proxy/temp-wrapper.sh",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-ed25519 AAA... pulse-sensor-proxy`
+	if !strings.HasPrefix(line, expectedStart) && line != expectedStart {
+		t.Errorf("unexpected authorized_key line: %s", line)
+	}
+
+	// Test empty config
+	proxyEmpty := &Proxy{config: &Config{}}
+	_, err = proxyEmpty.buildAuthorizedKey(pubKey)
+	if err == nil {
+		t.Error("expected error for empty allowed subnets")
+	}
+}
+
+func TestGetPublicKeyFrom(t *testing.T) {
+	tmpDir := t.TempDir()
+	pubKeyPath := filepath.Join(tmpDir, "id_ed25519.pub")
+
+	content := "ssh-ed25519 AAA... comment"
+	if err := os.WriteFile(pubKeyPath, []byte(content), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	// Test getPublicKeyFrom
+	proxy := &Proxy{}
+	pub, err := proxy.getPublicKeyFrom(tmpDir)
+	if err != nil {
+		t.Fatalf("getPublicKeyFrom failed: %v", err)
+	}
+
+	expected := "ssh-ed25519 AAA... comment"
+	if pub != expected {
+		t.Errorf("expected %s, got %s", expected, pub)
+	}
+
+	// Test missing
+	_, err = proxy.getPublicKeyFrom(filepath.Join(tmpDir, "missing"))
+	if err == nil {
+		t.Error("expected error for missing key")
+	}
+}
+
+func TestHandleHostKeyEnsureError(t *testing.T) {
+	proxy := &Proxy{
+		metrics: NewProxyMetrics("test"),
+	}
+
+	// Normal error
+	err := errors.New("normal error")
+	res := proxy.handleHostKeyEnsureError("node1", err)
+	if res != err {
+		t.Errorf("expected original error, got %v", res)
+	}
+
+	// HostKeyChangeError
+	changeErr := &knownhosts.HostKeyChangeError{
+		Host:     "node1",
+		Existing: "old",
+		Provided: "new",
+	}
+
+	res = proxy.handleHostKeyEnsureError("node1", changeErr)
+	if res != changeErr {
+		t.Errorf("expected original error, got %v", res)
+	}
+	// Verify metrics if possible, or just standard output log
+}
+
+func TestEnsureHostKeyFromProxmox(t *testing.T) {
+	// Mock isProxmoxHost
+	origLookPath := execLookPath
+	origStat := osStat
+	defer func() {
+		execLookPath = origLookPath
+		osStat = origStat
+	}()
+
+	// Case 1: Not Proxmox host
+	execLookPath = func(file string) (string, error) {
+		return "", fmt.Errorf("not found")
+	}
+	osStat = func(name string) (os.FileInfo, error) {
+		return nil, os.ErrNotExist
+	}
+
+	proxy := &Proxy{}
+	err := proxy.ensureHostKeyFromProxmox(context.Background(), "node1")
+	if err == nil || err.Error() != "not running on Proxmox host" {
+		t.Errorf("expected 'not running on Proxmox host' error, got %v", err)
+	}
+
+	// Case 2: Proxmox host, key found
+	execLookPath = func(file string) (string, error) {
+		if file == "pvecm" {
+			return "/bin/pvecm", nil
+		}
+		return "", fmt.Errorf("not found")
+	}
+
+	tmpDir := t.TempDir()
+	hostsFile := filepath.Join(tmpDir, "known_hosts")
+	origPath := proxmoxClusterKnownHostsPath
+	defer func() { proxmoxClusterKnownHostsPath = origPath }()
+	proxmoxClusterKnownHostsPath = hostsFile
+
+	os.WriteFile(hostsFile, []byte("node1 ssh-ed25519 AAA..."), 0644)
+
+	proxy = &Proxy{
+		knownHosts: &mockKnownHosts{},
+		config:     &Config{RequireProxmoxHostkeys: true},
+	}
+
+	err = proxy.ensureHostKeyFromProxmox(context.Background(), "node1")
+	if err != nil {
+		t.Errorf("ensureHostKeyFromProxmox failed: %v", err)
+	}
+
+	// Case 3: Proxmox host, key NOT found
+	// knownHosts EnsureWithEntries (mock) returns nil, but loadProxmoxHostKeys will verify if key exists in file
+
+	// loadProxmoxHostKeys returns error if key not found
+	err = proxy.ensureHostKeyFromProxmox(context.Background(), "unknown")
+	if err == nil {
+		t.Error("expected error for unknown host")
+	}
+}
+
+func TestPushSSHKeyFrom(t *testing.T) {
+	tmpDir := t.TempDir()
+	os.WriteFile(filepath.Join(tmpDir, "id_ed25519.pub"), []byte("ssh-ed25519 AAA..."), 0644)
+
+	proxy := &Proxy{
+		sshKeyPath: tmpDir,
+		metrics:    NewProxyMetrics("test"),
+		knownHosts: &mockKnownHosts{},
+		config:     &Config{AllowedSourceSubnets: []string{"10.0.0.0/24"}},
+	}
+
+	origExec := execCommandFunc
+	defer func() { execCommandFunc = origExec }()
+
+	// Case 1: Success (key already present)
+	execCommandFunc = func(name string, arg ...string) *exec.Cmd {
+		args := strings.Join(arg, " ")
+		if strings.Contains(args, "grep -F") {
+			return mockExecCommand("from=...,ssh-ed25519 AAA...") // Found
+		}
+		if strings.Contains(args, "stage temperature wrapper") {
+			return mockExecCommand("") // Wrapper install success
+		}
+		return mockExecCommand("")
+	}
+
+	if err := proxy.pushSSHKeyFrom("node1", tmpDir); err != nil {
+		t.Errorf("pushSSHKeyFrom failed (present): %v", err)
+	}
+
+	// Case 2: Success (key added)
+	execCommandFunc = func(name string, arg ...string) *exec.Cmd {
+		args := strings.Join(arg, " ")
+		if strings.Contains(args, "grep -F") {
+			return mockExecCommandWithExitCode("", 1) // Not found
+		}
+		return mockExecCommand("")
+	}
+
+	if err := proxy.pushSSHKeyFrom("node1", tmpDir); err != nil {
+		t.Errorf("pushSSHKeyFrom failed (added): %v", err)
+	}
+
+	// Case 3: Error key not found
+	if err := proxy.pushSSHKeyFrom("node1", filepath.Join(tmpDir, "missing")); err == nil {
+		t.Error("expected error for missing key dir")
+	}
+
+	// Case 4: ensureHostKey error (mock knownHosts.Ensure failure?)
+	// mockKnownHosts returns nil.
+}
+
+func TestSSHConnection(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	proxy := &Proxy{
+		sshKeyPath:        tmpDir,
+		metrics:           NewProxyMetrics("test"),
+		knownHosts:        &mockKnownHosts{},
+		maxSSHOutputBytes: 1024,
+		config:            &Config{},
+	}
+
+	origExec := execCommandFunc
+	origExecCtx := execCommandContextFunc
+	defer func() {
+		execCommandFunc = origExec
+		execCommandContextFunc = origExecCtx
+	}()
+
+	mockCmd := func(name string, arg ...string) *exec.Cmd {
+		if name == "sh" { // execCommandWithLimits uses "sh -c"
+			return mockExecCommand("{}")
+		}
+		return mockExecCommand("")
+	}
+
+	execCommandFunc = mockCmd
+	execCommandContextFunc = func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		return mockCmd(name, arg...)
+	}
+
+	if err := proxy.testSSHConnection("node1"); err != nil {
+		t.Errorf("testSSHConnection failed: %v", err)
+	}
+
+	// Failure
+	mockCmdFail := func(name string, arg ...string) *exec.Cmd {
+		if name == "sh" {
+			return mockExecCommandWithExitCode("failed", 1)
+		}
+		return mockExecCommand("")
+	}
+	execCommandContextFunc = func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+		return mockCmdFail(name, arg...)
+	}
+	execCommandFunc = mockCmdFail
+
+	if err := proxy.testSSHConnection("node1"); err == nil {
+		t.Error("expected error for connection failure")
+	}
+}