e2e: Add node.go to abstract node.SSH() and node.Reboot().

All nodes will share a global ssh client, which can be initialized
during the start of the tests by giving the path to the private key.

Author: Yifan Gu

e2e/main_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/crypto/ssh"
+
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/pkg/api/v1"

e2e/node.go

@@ -0,0 +1,49 @@
+package e2e
+
+import (
+	"fmt"
+
+	"k8s.io/client-go/pkg/api/v1"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type Node struct {
+	apiNode *v1.Node
+}
+
+func NewNode(n *v1.Node) *Node {
+	return &Node{apiNode: n}
+}
+
+func (n *Node) ExternalIP() string {
+	var host string
+	for _, addr := range n.apiNode.Status.Addresses {
+		if addr.Type == v1.NodeExternalIP {
+			host = addr.Address
+			break
+		}
+	}
+	return host
+}
+
+func (n *Node) SSH(cmd string) (stdout, stderr []byte, err error) {
+	host := n.ExternalIP()
+	if host == "" {
+		return nil, nil, fmt.Errorf("cannot find external IP for node %q", n.apiNode.Name)
+	}
+	return sshClient.SSH(host, cmd)
+}
+
+func (n *Node) Reboot() error {
+	stdout, stderr, err := n.SSH("sudo reboot")
+	if _, ok := err.(*ssh.ExitMissingError); ok {
+		// A terminated session is perfectly normal during reboot.
+		err = nil
+	}
+
+	if err != nil {
+		return fmt.Errorf("issuing reboot command failed\nstdout:%s\nstderr:%s", stdout, stderr)
+	}
+	return nil
+}

e2e/ssh_client.go

@@ -11,10 +11,17 @@ import (
 	"golang.org/x/crypto/ssh/agent"
 )
 
+// The global SSH client used by the tests.
+var sshClient *SSHClient
+
 type SSHClient struct {
 	*ssh.ClientConfig
 }
 
+func InitSSHClient(keypath string) {
+	sshClient = NewSSHClientOrDie(keypath)
+}
+
 // NewSSHClientOrDie tries to create an ssh client.
 // If $SSH_AUTH_SOCK is set, the use the ssh agent to create the client,
 // otherwise read the private key directly.