I wrote an SSH client in Go and I would like to write some tests. The problem is that I've never really written proper unit tests before, and most tutorials seem to focus on writing tests for a function that adds two numbers or some other toy problem. I've read about mocking, using interfaces, and other techniques, but I'm having trouble applying them. Also, my client is going to be used concurrently to allow fast configuration of multiple devices at a time. Not sure if that would change the way I write my tests or would add additional tests. Any help is appreciated.
Here is my code. Basically, a Device has 4 main functions: Connect, Send, Output/Err and Close for connecting to a device, sending it a set of configuration commands, capturing the output of the session, and closing the client, respectively.
package device
import (
"bufio"
"fmt"
"golang.org/x/crypto/ssh"
"io"
"net"
"time"
)
// A Device represents a remote network device.
type Device struct {
Host string // the device's hostname or IP address
client *ssh.Client // the client connection
session *ssh.Session // the connection to the remote shell
stdin io.WriteCloser // the remote shell's standard input
stdout io.Reader // the remote shell's standard output
stderr io.Reader // the remote shell's standard error
}
// Connect establishes an SSH connection to a device and sets up the session IO.
func (d *Device) Connect(user, password string) error {
// Create a client connection
client, err := ssh.Dial("tcp", net.JoinHostPort(d.Host, "22"), configureClient(user, password))
if err != nil {
return err
}
d.client = client
// Create a session
session, err := client.NewSession()
if err != nil {
return err
}
d.session = session
return nil
}
// configureClient sets up the client configuration for login
func configureClient(user, password string) *ssh.ClientConfig {
var sshConfig ssh.Config
sshConfig.SetDefaults()
sshConfig.Ciphers = append(sshConfig.Ciphers, "aes128-cbc", "aes256-cbc", "3des-cbc", "des-cbc", "aes192-cbc")
config := &ssh.ClientConfig{
Config: sshConfig,
User: user,
Auth: []ssh.AuthMethod{ssh.Password(password)},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
Timeout: time.Second * 5,
}
return config
}
// setupIO creates the pipes connected to the remote shell's standard input, output, and error
func (d *Device) setupIO() error {
// Setup standard input pipe
stdin, err := d.session.StdinPipe()
if err != nil {
return err
}
d.stdin = stdin
// Setup standard output pipe
stdout, err := d.session.StdoutPipe()
if err != nil {
return err
}
d.stdout = stdout
// Setup standard error pipe
stderr, err := d.session.StderrPipe()
if err != nil {
return err
}
d.stderr = stderr
return nil
}
// Send sends cmd(s) to the device's standard input. A device only accepts one call
// to Send, as it closes the session and its standard input pipe.
func (d *Device) Send(cmds ...string) error {
if d.session == nil {
return fmt.Errorf("device: session is closed")
}
defer d.session.Close()
// Start the shell
if err := d.startShell(); err != nil {
return err
}
// Send commands
for _, cmd := range cmds {
if _, err := d.stdin.Write([]byte(cmd + "")); err != nil {
return err
}
}
defer d.stdin.Close()
// Wait for the commands to exit
d.session.Wait()
return nil
}
// startShell requests a pseudo terminal (VT100) and starts the remote shell.
func (d *Device) startShell() error {
modes := ssh.TerminalModes{
ssh.ECHO: 0, // disable echoing
ssh.OCRNL: 0,
ssh.TTY_OP_ISPEED: 14400,
ssh.TTY_OP_OSPEED: 14400,
}
err := d.session.RequestPty("vt100", 0, 0, modes)
if err != nil {
return err
}
if err := d.session.Shell(); err != nil {
return err
}
return nil
}
// Output returns the remote device's standard output output.
func (d *Device) Output() ([]string, error) {
return readPipe(d.stdout)
}
// Err returns the remote device's standard error output.
func (d *Device) Err() ([]string, error) {
return readPipe(d.stdout)
}
// reapPipe reads an io.Reader line by line
func readPipe(r io.Reader) ([]string, error) {
var lines []string
scanner := bufio.NewScanner(r)
for scanner.Scan() {
lines = append(lines, scanner.Text())
}
if err := scanner.Err(); err != nil {
return nil, err
}
return lines, nil
}
// Close closes the client connection.
func (d *Device) Close() error {
return d.client.Close()
}
// String returns the string representation of a `Device`.
func (d *Device) String() string {
return fmt.Sprintf("%s", d.Host)
}
You make a good point about unit test tutorials nearly always being toy problems (why is it always Fibonacci?), when what we have is databases and http servers. The big realization that helped me is that you can only unit test things where you can control the input and output of the unit. configureClient or readPipe (give it a strings.Reader) would be good candidates. Start there.
Anything that leaves your program by talking directly to the disk, the network, stdout, etc, like the Connect method you would consider part of the external interface of your program. You don't unit test those. You integration test them.
Change Device to be an interface rather than a struct, and make a MockDevice that implements it. The real device is now maybe SSHDevice. You can unit test the rest of your program (which uses Device interface) by inserting a MockDevice, to isolate yourself from the network.
The SSHDevice will get tested in your integration tests. Start a real ssh server (maybe a test one you write in Go using crypto/ssh package, but any sshd would work). Start your program with an SSHDevice, make them talk to each other, and check outputs. You'll be using the os/exec package a lot. Integration tests are even more fun to write than unit tests!