launchpad.net/~rogpeppe/juju-core/500-errgo-fix@v0.0.0-20140213181702-000000002356/cmd/juju/ssh_test.go (about) 1 // Copyright 2012, 2013 Canonical Ltd. 2 // Licensed under the AGPLv3, see LICENCE file for details. 3 4 package main 5 6 import ( 7 "bytes" 8 "fmt" 9 "net/url" 10 "os" 11 "path/filepath" 12 13 gc "launchpad.net/gocheck" 14 15 "launchpad.net/juju-core/charm" 16 "launchpad.net/juju-core/cmd" 17 "launchpad.net/juju-core/instance" 18 "launchpad.net/juju-core/juju/testing" 19 "launchpad.net/juju-core/state" 20 coretesting "launchpad.net/juju-core/testing" 21 ) 22 23 var _ = gc.Suite(&SSHSuite{}) 24 25 type SSHSuite struct { 26 SSHCommonSuite 27 } 28 29 type SSHCommonSuite struct { 30 testing.JujuConnSuite 31 bin string 32 } 33 34 // fakecommand outputs its arguments to stdout for verification 35 var fakecommand = `#!/bin/bash 36 37 echo $@ | tee $0.args 38 ` 39 40 func (s *SSHCommonSuite) SetUpTest(c *gc.C) { 41 s.JujuConnSuite.SetUpTest(c) 42 43 s.bin = c.MkDir() 44 s.PatchEnvironment("PATH", s.bin+":"+os.Getenv("PATH")) 45 for _, name := range []string{"ssh", "scp"} { 46 f, err := os.OpenFile(filepath.Join(s.bin, name), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0777) 47 c.Assert(err, gc.IsNil) 48 _, err = f.Write([]byte(fakecommand)) 49 c.Assert(err, gc.IsNil) 50 err = f.Close() 51 c.Assert(err, gc.IsNil) 52 } 53 } 54 55 const ( 56 commonArgs = `-o StrictHostKeyChecking no -o PasswordAuthentication no ` 57 sshArgs = commonArgs + `-t ` 58 ) 59 60 var sshTests = []struct { 61 args []string 62 result string 63 }{ 64 { 65 []string{"ssh", "0"}, 66 sshArgs + "ubuntu@dummyenv-0.dns\n", 67 }, 68 // juju ssh 0 'uname -a' 69 { 70 []string{"ssh", "0", "uname -a"}, 71 sshArgs + "ubuntu@dummyenv-0.dns -- uname -a\n", 72 }, 73 // juju ssh 0 -- uname -a 74 { 75 []string{"ssh", "0", "--", "uname", "-a"}, 76 sshArgs + "ubuntu@dummyenv-0.dns -- uname -a\n", 77 }, 78 // juju ssh 0 uname -a 79 { 80 []string{"ssh", "0", "uname", "-a"}, 81 sshArgs + "ubuntu@dummyenv-0.dns -- uname -a\n", 82 }, 83 { 84 []string{"ssh", "mysql/0"}, 85 sshArgs + "ubuntu@dummyenv-0.dns\n", 86 }, 87 { 88 []string{"ssh", "mongodb/1"}, 89 sshArgs + "ubuntu@dummyenv-2.dns\n", 90 }, 91 } 92 93 func (s *SSHSuite) TestSSHCommand(c *gc.C) { 94 m := s.makeMachines(3, c, true) 95 ch := coretesting.Charms.Dir("dummy") 96 curl := charm.MustParseURL( 97 fmt.Sprintf("local:quantal/%s-%d", ch.Meta().Name, ch.Revision()), 98 ) 99 bundleURL, err := url.Parse("http://bundles.testing.invalid/dummy-1") 100 c.Assert(err, gc.IsNil) 101 dummy, err := s.State.AddCharm(ch, curl, bundleURL, "dummy-1-sha256") 102 c.Assert(err, gc.IsNil) 103 srv := s.AddTestingService(c, "mysql", dummy) 104 s.addUnit(srv, m[0], c) 105 106 srv = s.AddTestingService(c, "mongodb", dummy) 107 s.addUnit(srv, m[1], c) 108 s.addUnit(srv, m[2], c) 109 110 for _, t := range sshTests { 111 c.Logf("testing juju ssh %s", t.args) 112 ctx := coretesting.Context(c) 113 jujucmd := cmd.NewSuperCommand(cmd.SuperCommandParams{}) 114 jujucmd.Register(&SSHCommand{}) 115 116 code := cmd.Main(jujucmd, ctx, t.args) 117 c.Check(code, gc.Equals, 0) 118 c.Check(ctx.Stderr.(*bytes.Buffer).String(), gc.Equals, "") 119 c.Check(ctx.Stdout.(*bytes.Buffer).String(), gc.Equals, t.result) 120 } 121 } 122 123 type callbackAttemptStarter struct { 124 next func() bool 125 } 126 127 func (s *callbackAttemptStarter) Start() attempt { 128 return callbackAttempt{next: s.next} 129 } 130 131 type callbackAttempt struct { 132 next func() bool 133 } 134 135 func (a callbackAttempt) Next() bool { 136 return a.next() 137 } 138 139 func (s *SSHSuite) TestSSHCommandHostAddressRetry(c *gc.C) { 140 m := s.makeMachines(1, c, false) 141 ctx := coretesting.Context(c) 142 jujucmd := cmd.NewSuperCommand(cmd.SuperCommandParams{}) 143 jujucmd.Register(&SSHCommand{}) 144 145 var called int 146 next := func() bool { 147 called++ 148 return called < 2 149 } 150 attemptStarter := &callbackAttemptStarter{next: next} 151 s.PatchValue(&sshHostFromTargetAttemptStrategy, attemptStarter) 152 153 // Ensure that the ssh command waits for a public address, or the attempt 154 // strategy's Done method returns false. 155 code := cmd.Main(jujucmd, ctx, []string{"ssh", "0"}) 156 c.Check(code, gc.Equals, 1) 157 c.Assert(called, gc.Equals, 2) 158 called = 0 159 attemptStarter.next = func() bool { 160 called++ 161 s.setAddress(m[0], c) 162 return false 163 } 164 code = cmd.Main(jujucmd, ctx, []string{"ssh", "0"}) 165 c.Check(code, gc.Equals, 0) 166 c.Assert(called, gc.Equals, 1) 167 } 168 169 func (s *SSHCommonSuite) setAddress(m *state.Machine, c *gc.C) { 170 addr := instance.NewAddress(fmt.Sprintf("dummyenv-%s.dns", m.Id())) 171 addr.NetworkScope = instance.NetworkPublic 172 err := m.SetAddresses([]instance.Address{addr}) 173 c.Assert(err, gc.IsNil) 174 } 175 176 func (s *SSHCommonSuite) makeMachines(n int, c *gc.C, setAddress bool) []*state.Machine { 177 var machines = make([]*state.Machine, n) 178 for i := 0; i < n; i++ { 179 m, err := s.State.AddMachine("quantal", state.JobHostUnits) 180 c.Assert(err, gc.IsNil) 181 if setAddress { 182 s.setAddress(m, c) 183 } 184 // must set an instance id as the ssh command uses that as a signal the 185 // machine has been provisioned 186 inst, md := testing.AssertStartInstance(c, s.Conn.Environ, m.Id()) 187 c.Assert(m.SetProvisioned(inst.Id(), "fake_nonce", md), gc.IsNil) 188 machines[i] = m 189 } 190 return machines 191 } 192 193 func (s *SSHCommonSuite) addUnit(srv *state.Service, m *state.Machine, c *gc.C) { 194 u, err := srv.AddUnit() 195 c.Assert(err, gc.IsNil) 196 err = u.AssignToMachine(m) 197 c.Assert(err, gc.IsNil) 198 // fudge unit.SetPublicAddress 199 id, err := m.InstanceId() 200 c.Assert(err, gc.IsNil) 201 insts, err := s.Conn.Environ.Instances([]instance.Id{id}) 202 c.Assert(err, gc.IsNil) 203 addr, err := insts[0].WaitDNSName() 204 c.Assert(err, gc.IsNil) 205 err = u.SetPublicAddress(addr) 206 c.Assert(err, gc.IsNil) 207 }