From 2c7e632925a1236589778855b45c89fb5ede8c4a Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Fri, 11 Dec 2015 17:56:00 -0800
Subject: [PATCH] Amend rotation tests to assert old keys are removed after
 rotation.

Signed-off-by: Ying Li <ying.li@docker.com>
---
 client/client.go      |  6 +-----
 client/client_test.go | 17 +++++++++++++----
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/client/client.go b/client/client.go
index c797c52a7f..c498a62ddb 100644
--- a/client/client.go
+++ b/client/client.go
@@ -653,11 +653,7 @@ func (r *NotaryRepository) RotateKey(role string, serverManagesKey bool) error {
 		return err
 	}
 
-	err = r.rootFileKeyChange(role, changelist.ActionCreate, pubKey)
-	if err != nil {
-		return err
-	}
-	return nil
+	return r.rootFileKeyChange(role, changelist.ActionCreate, pubKey)
 }
 
 func (r *NotaryRepository) rootFileKeyChange(role, action string, key data.PublicKey) error {
diff --git a/client/client_test.go b/client/client_test.go
index b90bc5869c..9bbba350be 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -1055,11 +1055,10 @@ func TestRotateKeyInvalidRole(t *testing.T) {
 func assertRotationSuccessful(t *testing.T, repo *NotaryRepository,
 	keysToRotate map[string]bool) {
 
-	oldKeyIDs := make(map[string]string)
+	oldKeyIDs := make(map[string][]string)
 	for role := range keysToRotate {
 		keyIDs := repo.tufRepo.Root.Signed.Roles[role].KeyIDs
-		assert.Len(t, keyIDs, 1)
-		oldKeyIDs[role] = keyIDs[0]
+		oldKeyIDs[role] = keyIDs
 	}
 
 	// Do rotation
@@ -1078,7 +1077,17 @@ func assertRotationSuccessful(t *testing.T, repo *NotaryRepository,
 	for role, isRemoteKey := range keysToRotate {
 		keyIDs := repo.tufRepo.Root.Signed.Roles[role].KeyIDs
 		assert.Len(t, keyIDs, 1)
-		assert.NotEqual(t, oldKeyIDs[role], keyIDs[0])
+
+		// the new key is not the same as any of the old keys, and the
+		// old keys have been removed not just from the TUF file, but
+		// from the cryptoservice
+		for _, oldKeyID := range oldKeyIDs[role] {
+			assert.NotEqual(t, oldKeyID, keyIDs[0])
+			_, _, err := repo.CryptoService.GetPrivateKey(oldKeyID)
+			assert.Error(t, err)
+		}
+
+		// the new key is present in the cryptoservice, or not present if remote
 		key, _, err := repo.CryptoService.GetPrivateKey(keyIDs[0])
 		if isRemoteKey {
 			assert.Error(t, err)