From 17173efbe00d422392f91c672350266314808b28 Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <aaron.lehmann@docker.com>
Date: Tue, 6 Jun 2017 16:28:40 +0300
Subject: [PATCH] integration-cli: Replace sleeps with polling in swarm
 lock/unlock tests

This will hopefully make the tests more robust by replacing a fixed 3s
sleep with a polling loop that looks at whether the key PEM file is
encrypted or not.

Signed-off-by: Aaron Lehmann <aaron.lehmann@docker.com>
---
 integration-cli/docker_cli_swarm_test.go | 39 +++++++++++++-----------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/integration-cli/docker_cli_swarm_test.go b/integration-cli/docker_cli_swarm_test.go
index 086a6632d2..a7d4a60aa8 100644
--- a/integration-cli/docker_cli_swarm_test.go
+++ b/integration-cli/docker_cli_swarm_test.go
@@ -4,7 +4,9 @@ package main
 
 import (
 	"bytes"
+	"crypto/x509"
 	"encoding/json"
+	"encoding/pem"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -992,32 +994,35 @@ func getNodeStatus(c *check.C, d *daemon.Swarm) swarm.LocalNodeState {
 	return info.LocalNodeState
 }
 
-func checkSwarmLockedToUnlocked(c *check.C, d *daemon.Swarm, unlockKey string) {
-	d.Restart(c)
-	status := getNodeStatus(c, d)
-	if status == swarm.LocalNodeStateLocked {
-		// it must not have updated to be unlocked in time - unlock, wait 3 seconds, and try again
-		cmd := d.Command("swarm", "unlock")
-		cmd.Stdin = bytes.NewBufferString(unlockKey)
-		icmd.RunCmd(cmd).Assert(c, icmd.Success)
+func checkKeyIsEncrypted(d *daemon.Swarm) func(*check.C) (interface{}, check.CommentInterface) {
+	return func(c *check.C) (interface{}, check.CommentInterface) {
+		keyBytes, err := ioutil.ReadFile(filepath.Join(d.Folder, "root", "swarm", "certificates", "swarm-node.key"))
+		if err != nil {
+			return fmt.Errorf("error reading key: %v", err), nil
+		}
 
-		c.Assert(getNodeStatus(c, d), checker.Equals, swarm.LocalNodeStateActive)
+		keyBlock, _ := pem.Decode(keyBytes)
+		if keyBlock == nil {
+			return fmt.Errorf("invalid PEM-encoded private key"), nil
+		}
 
-		time.Sleep(3 * time.Second)
-		d.Restart(c)
+		return x509.IsEncryptedPEMBlock(keyBlock), nil
 	}
+}
 
+func checkSwarmLockedToUnlocked(c *check.C, d *daemon.Swarm, unlockKey string) {
+	// Wait for the PEM file to become unencrypted
+	waitAndAssert(c, defaultReconciliationTimeout, checkKeyIsEncrypted(d), checker.Equals, false)
+
+	d.Restart(c)
 	c.Assert(getNodeStatus(c, d), checker.Equals, swarm.LocalNodeStateActive)
 }
 
 func checkSwarmUnlockedToLocked(c *check.C, d *daemon.Swarm) {
+	// Wait for the PEM file to become encrypted
+	waitAndAssert(c, defaultReconciliationTimeout, checkKeyIsEncrypted(d), checker.Equals, true)
+
 	d.Restart(c)
-	status := getNodeStatus(c, d)
-	if status == swarm.LocalNodeStateActive {
-		// it must not have updated to be unlocked in time - wait 3 seconds, and try again
-		time.Sleep(3 * time.Second)
-		d.Restart(c)
-	}
 	c.Assert(getNodeStatus(c, d), checker.Equals, swarm.LocalNodeStateLocked)
 }