github.com/mattyw/juju@v0.0.0-20140610034352-732aecd63861/worker/authenticationworker/worker_test.go (about) 1 // Copyright 2013 Canonical Ltd. 2 // Licensed under the AGPLv3, see LICENCE file for details. 3 4 package authenticationworker_test 5 6 import ( 7 "strings" 8 stdtesting "testing" 9 "time" 10 11 gc "launchpad.net/gocheck" 12 13 "github.com/juju/juju/agent" 14 jujutesting "github.com/juju/juju/juju/testing" 15 "github.com/juju/juju/state" 16 "github.com/juju/juju/state/api" 17 "github.com/juju/juju/state/api/keyupdater" 18 coretesting "github.com/juju/juju/testing" 19 "github.com/juju/juju/utils/ssh" 20 sshtesting "github.com/juju/juju/utils/ssh/testing" 21 "github.com/juju/juju/worker" 22 "github.com/juju/juju/worker/authenticationworker" 23 ) 24 25 // worstCase is used for timeouts when timing out 26 // will fail the test. Raising this value should 27 // not affect the overall running time of the tests 28 // unless they fail. 29 const worstCase = 5 * time.Second 30 31 func TestAll(t *stdtesting.T) { 32 coretesting.MgoTestPackage(t) 33 } 34 35 var _ = gc.Suite(&workerSuite{}) 36 37 type workerSuite struct { 38 jujutesting.JujuConnSuite 39 stateMachine *state.Machine 40 machine *state.Machine 41 keyupdaterApi *keyupdater.State 42 43 existingEnvKey string 44 existingKeys []string 45 } 46 47 func (s *workerSuite) SetUpTest(c *gc.C) { 48 s.JujuConnSuite.SetUpTest(c) 49 // Default ssh user is currently "ubuntu". 50 c.Assert(authenticationworker.SSHUser, gc.Equals, "ubuntu") 51 // Set the ssh user to empty (the current user) as required by the test infrastructure. 52 s.PatchValue(&authenticationworker.SSHUser, "") 53 54 // Replace the default dummy key in the test environment with a valid one. 55 // This will be added to the ssh authorised keys when the agent starts. 56 s.setAuthorisedKeys(c, sshtesting.ValidKeyOne.Key+" firstuser@host") 57 // Record the existing key with its prefix for testing later. 58 s.existingEnvKey = sshtesting.ValidKeyOne.Key + " Juju:firstuser@host" 59 60 // Set up an existing key (which is not in the environment) in the ssh authorised_keys file. 61 s.existingKeys = []string{sshtesting.ValidKeyTwo.Key + " existinguser@host"} 62 err := ssh.AddKeys(authenticationworker.SSHUser, s.existingKeys...) 63 c.Assert(err, gc.IsNil) 64 65 var apiRoot *api.State 66 apiRoot, s.machine = s.OpenAPIAsNewMachine(c) 67 c.Assert(apiRoot, gc.NotNil) 68 s.keyupdaterApi = apiRoot.KeyUpdater() 69 c.Assert(s.keyupdaterApi, gc.NotNil) 70 } 71 72 func stop(c *gc.C, w worker.Worker) { 73 c.Assert(worker.Stop(w), gc.IsNil) 74 } 75 76 type mockConfig struct { 77 agent.Config 78 c *gc.C 79 tag string 80 } 81 82 func (mock *mockConfig) Tag() string { 83 return mock.tag 84 } 85 86 func agentConfig(c *gc.C, tag string) *mockConfig { 87 return &mockConfig{c: c, tag: tag} 88 } 89 90 func (s *workerSuite) setAuthorisedKeys(c *gc.C, keys ...string) { 91 keyStr := strings.Join(keys, "\n") 92 err := s.BackingState.UpdateEnvironConfig(map[string]interface{}{"authorized-keys": keyStr}, nil, nil) 93 c.Assert(err, gc.IsNil) 94 s.BackingState.StartSync() 95 } 96 97 func (s *workerSuite) waitSSHKeys(c *gc.C, expected []string) { 98 timeout := time.After(worstCase) 99 for { 100 select { 101 case <-timeout: 102 c.Fatalf("timeout while waiting for authoirsed ssh keys to change") 103 case <-time.After(coretesting.ShortWait): 104 keys, err := ssh.ListKeys(authenticationworker.SSHUser, ssh.FullKeys) 105 c.Assert(err, gc.IsNil) 106 keysStr := strings.Join(keys, "\n") 107 expectedStr := strings.Join(expected, "\n") 108 if expectedStr != keysStr { 109 continue 110 } 111 return 112 } 113 } 114 } 115 116 func (s *workerSuite) TestKeyUpdateRetainsExisting(c *gc.C) { 117 authWorker := authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 118 defer stop(c, authWorker) 119 120 newKey := sshtesting.ValidKeyThree.Key + " user@host" 121 s.setAuthorisedKeys(c, newKey) 122 newKeyWithCommentPrefix := sshtesting.ValidKeyThree.Key + " Juju:user@host" 123 s.waitSSHKeys(c, append(s.existingKeys, newKeyWithCommentPrefix)) 124 } 125 126 func (s *workerSuite) TestNewKeysInJujuAreSavedOnStartup(c *gc.C) { 127 newKey := sshtesting.ValidKeyThree.Key + " user@host" 128 s.setAuthorisedKeys(c, newKey) 129 130 authWorker := authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 131 defer stop(c, authWorker) 132 133 newKeyWithCommentPrefix := sshtesting.ValidKeyThree.Key + " Juju:user@host" 134 s.waitSSHKeys(c, append(s.existingKeys, newKeyWithCommentPrefix)) 135 } 136 137 func (s *workerSuite) TestDeleteKey(c *gc.C) { 138 authWorker := authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 139 defer stop(c, authWorker) 140 141 // Add another key 142 anotherKey := sshtesting.ValidKeyThree.Key + " another@host" 143 s.setAuthorisedKeys(c, s.existingEnvKey, anotherKey) 144 anotherKeyWithCommentPrefix := sshtesting.ValidKeyThree.Key + " Juju:another@host" 145 s.waitSSHKeys(c, append(s.existingKeys, s.existingEnvKey, anotherKeyWithCommentPrefix)) 146 147 // Delete the original key and check anotherKey plus the existing keys remain. 148 s.setAuthorisedKeys(c, anotherKey) 149 s.waitSSHKeys(c, append(s.existingKeys, anotherKeyWithCommentPrefix)) 150 } 151 152 func (s *workerSuite) TestMultipleChanges(c *gc.C) { 153 authWorker := authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 154 defer stop(c, authWorker) 155 s.waitSSHKeys(c, append(s.existingKeys, s.existingEnvKey)) 156 157 // Perform a set to add a key and delete a key. 158 // added: key 3 159 // deleted: key 1 (existing env key) 160 s.setAuthorisedKeys(c, sshtesting.ValidKeyThree.Key+" yetanother@host") 161 yetAnotherKeyWithComment := sshtesting.ValidKeyThree.Key + " Juju:yetanother@host" 162 s.waitSSHKeys(c, append(s.existingKeys, yetAnotherKeyWithComment)) 163 } 164 165 func (s *workerSuite) TestWorkerRestart(c *gc.C) { 166 authWorker := authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 167 defer stop(c, authWorker) 168 s.waitSSHKeys(c, append(s.existingKeys, s.existingEnvKey)) 169 170 // Stop the worker and delete and add keys from the environment while it is down. 171 // added: key 3 172 // deleted: key 1 (existing env key) 173 stop(c, authWorker) 174 s.setAuthorisedKeys(c, sshtesting.ValidKeyThree.Key+" yetanother@host") 175 176 // Restart the worker and check that the ssh auth keys are as expected. 177 authWorker = authenticationworker.NewWorker(s.keyupdaterApi, agentConfig(c, s.machine.Tag())) 178 defer stop(c, authWorker) 179 180 yetAnotherKeyWithCommentPrefix := sshtesting.ValidKeyThree.Key + " Juju:yetanother@host" 181 s.waitSSHKeys(c, append(s.existingKeys, yetAnotherKeyWithCommentPrefix)) 182 }