e2e: Add simple cluster type to make tests migration easier.

Calling GetCluster() in each test can returns a Cluster type that
stores the list of master and worker nodes, which provide SSH() and
Reboot() interfaces.

Author: Yifan Gu

e2e/main_test.go

@@ -7,8 +7,6 @@ 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

@@ -3,22 +3,27 @@ package e2e
 import (
 	"fmt"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/pkg/api/v1"
 
 	"golang.org/x/crypto/ssh"
 )
 
+const (
+	LabelNodeRoleMaster = "node-role.kubernetes.io/master"
+)
+
 type Node struct {
-	apiNode *v1.Node
+	*v1.Node
 }
 
-func NewNode(n *v1.Node) *Node {
-	return &Node{apiNode: n}
+func newNode(n *v1.Node) *Node {
+	return &Node{n}
 }
 
 func (n *Node) ExternalIP() string {
 	var host string
-	for _, addr := range n.apiNode.Status.Addresses {
+	for _, addr := range n.Status.Addresses {
 		if addr.Type == v1.NodeExternalIP {
 			host = addr.Address
 			break
@@ -30,7 +35,7 @@ func (n *Node) ExternalIP() string {
 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 nil, nil, fmt.Errorf("cannot find external IP for node %q", n.Name)
 	}
 	return sshClient.SSH(host, cmd)
 }
@@ -47,3 +52,35 @@ func (n *Node) Reboot() error {
 	}
 	return nil
 }
+
+// IsMaster returns true if the node's labels contains "node-role.kubernetes.io/master".
+func (n *Node) IsMaster() bool {
+	_, ok := n.Labels[LabelNodeRoleMaster]
+	return ok
+}
+
+// Cluster is a simple abstraction to make writing tests easier.
+type Cluster struct {
+	Masters []*Node
+	Workers []*Node
+}
+
+// GetCluster can be called in every test to return a *Cluster object.
+func GetCluster() (*Cluster, error) {
+	var c Cluster
+
+	nodelist, err := client.CoreV1().Nodes().List(metav1.ListOptions{})
+	if err != nil {
+		return nil, err
+	}
+
+	for _, n := range nodelist.Items {
+		nn := newNode(&n)
+		if nn.IsMaster() {
+			c.Masters = append(c.Masters, nn)
+		} else {
+			c.Workers = append(c.Workers, nn)
+		}
+	}
+	return &c, nil
+}

e2e/ssh_client.go

@@ -19,13 +19,13 @@ type SSHClient struct {
 }
 
 func InitSSHClient(keypath string) {
-	sshClient = NewSSHClientOrDie(keypath)
+	sshClient = newSSHClientOrDie(keypath)
 }
 
-// NewSSHClientOrDie tries to create an ssh client.
+// 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.
-func NewSSHClientOrDie(keypath string) *SSHClient {
+func newSSHClientOrDie(keypath string) *SSHClient {
 	var authMethod ssh.AuthMethod
 
 	sock := os.Getenv("SSH_AUTH_SOCK")