starssh/keepalive_test.go

package starssh

import (
	"context"
	"testing"
	"time"

	"golang.org/x/crypto/ssh"
)

func TestAutoKeepAliveTimeoutDoesNotDeadlock(t *testing.T) {
	oldSendKeepAliveRequest := sendKeepAliveRequest
	oldCloseSSHClient := closeSSHClient
	t.Cleanup(func() {
		sendKeepAliveRequest = oldSendKeepAliveRequest
		closeSSHClient = oldCloseSSHClient
	})

	sendKeepAliveRequest = func(ctx context.Context, client sshClientRequester) error {
		<-ctx.Done()
		return ctx.Err()
	}
	closeSSHClient = func(client sshClientRequester) error {
		return nil
	}

	client := &ssh.Client{}
	star := &StarSSH{
		LoginInfo: LoginInput{
			KeepAliveInterval: 10 * time.Millisecond,
			KeepAliveTimeout:  20 * time.Millisecond,
		},
	}
	star.setTransport(client, nil)
	star.startAutoKeepAlive()

	star.keepaliveMu.Lock()
	done := star.keepaliveDone
	star.keepaliveMu.Unlock()
	if done == nil {
		t.Fatal("keepalive did not start")
	}

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("keepalive goroutine did not exit after keepalive timeout")
	}

	if client := star.snapshotSSHClient(); client != nil {
		t.Fatal("ssh client should be closed after keepalive timeout")
	}
}