github.com/cloudbase/juju-core@v0.0.0-20140504232958-a7271ac7912f/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.PatchEnvPathPrepend(s.bin) 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 -t ` 58 ) 59 60 var sshTests = []struct { 61 about string 62 args []string 63 result string 64 }{ 65 { 66 "connect to machine 0", 67 []string{"ssh", "0"}, 68 sshArgs + "ubuntu@dummyenv-0.dns\n", 69 }, 70 { 71 "connect to machine 0 and pass extra arguments", 72 []string{"ssh", "0", "uname", "-a"}, 73 sshArgs + "ubuntu@dummyenv-0.dns uname -a\n", 74 }, 75 { 76 "connect to unit mysql/0", 77 []string{"ssh", "mysql/0"}, 78 sshArgs + "ubuntu@dummyenv-0.dns\n", 79 }, 80 { 81 "connect to unit mongodb/1 and pass extra arguments", 82 []string{"ssh", "mongodb/1", "ls", "/"}, 83 sshArgs + "ubuntu@dummyenv-2.dns ls /\n", 84 }, 85 } 86 87 func (s *SSHSuite) TestSSHCommand(c *gc.C) { 88 m := s.makeMachines(3, c, true) 89 ch := coretesting.Charms.Dir("dummy") 90 curl := charm.MustParseURL( 91 fmt.Sprintf("local:quantal/%s-%d", ch.Meta().Name, ch.Revision()), 92 ) 93 bundleURL, err := url.Parse("http://bundles.testing.invalid/dummy-1") 94 c.Assert(err, gc.IsNil) 95 dummy, err := s.State.AddCharm(ch, curl, bundleURL, "dummy-1-sha256") 96 c.Assert(err, gc.IsNil) 97 srv := s.AddTestingService(c, "mysql", dummy) 98 s.addUnit(srv, m[0], c) 99 100 srv = s.AddTestingService(c, "mongodb", dummy) 101 s.addUnit(srv, m[1], c) 102 s.addUnit(srv, m[2], c) 103 104 for i, t := range sshTests { 105 c.Logf("test %d: %s -> %s\n", i, t.about, t.args) 106 ctx := coretesting.Context(c) 107 jujucmd := cmd.NewSuperCommand(cmd.SuperCommandParams{}) 108 jujucmd.Register(&SSHCommand{}) 109 110 code := cmd.Main(jujucmd, ctx, t.args) 111 c.Check(code, gc.Equals, 0) 112 c.Check(ctx.Stderr.(*bytes.Buffer).String(), gc.Equals, "") 113 c.Check(ctx.Stdout.(*bytes.Buffer).String(), gc.Equals, t.result) 114 } 115 } 116 117 type callbackAttemptStarter struct { 118 next func() bool 119 } 120 121 func (s *callbackAttemptStarter) Start() attempt { 122 return callbackAttempt{next: s.next} 123 } 124 125 type callbackAttempt struct { 126 next func() bool 127 } 128 129 func (a callbackAttempt) Next() bool { 130 return a.next() 131 } 132 133 func (s *SSHSuite) TestSSHCommandHostAddressRetry(c *gc.C) { 134 m := s.makeMachines(1, c, false) 135 ctx := coretesting.Context(c) 136 jujucmd := cmd.NewSuperCommand(cmd.SuperCommandParams{}) 137 jujucmd.Register(&SSHCommand{}) 138 139 var called int 140 next := func() bool { 141 called++ 142 return called < 2 143 } 144 attemptStarter := &callbackAttemptStarter{next: next} 145 s.PatchValue(&sshHostFromTargetAttemptStrategy, attemptStarter) 146 147 // Ensure that the ssh command waits for a public address, or the attempt 148 // strategy's Done method returns false. 149 code := cmd.Main(jujucmd, ctx, []string{"ssh", "0"}) 150 c.Check(code, gc.Equals, 1) 151 c.Assert(called, gc.Equals, 2) 152 called = 0 153 attemptStarter.next = func() bool { 154 called++ 155 s.setAddress(m[0], c) 156 return false 157 } 158 code = cmd.Main(jujucmd, ctx, []string{"ssh", "0"}) 159 c.Check(code, gc.Equals, 0) 160 c.Assert(called, gc.Equals, 1) 161 } 162 163 func (s *SSHCommonSuite) setAddress(m *state.Machine, c *gc.C) { 164 addr := instance.NewAddress(fmt.Sprintf("dummyenv-%s.dns", m.Id())) 165 addr.NetworkScope = instance.NetworkPublic 166 err := m.SetAddresses([]instance.Address{addr}) 167 c.Assert(err, gc.IsNil) 168 } 169 170 func (s *SSHCommonSuite) makeMachines(n int, c *gc.C, setAddress bool) []*state.Machine { 171 var machines = make([]*state.Machine, n) 172 for i := 0; i < n; i++ { 173 m, err := s.State.AddMachine("quantal", state.JobHostUnits) 174 c.Assert(err, gc.IsNil) 175 if setAddress { 176 s.setAddress(m, c) 177 } 178 // must set an instance id as the ssh command uses that as a signal the 179 // machine has been provisioned 180 inst, md := testing.AssertStartInstance(c, s.Conn.Environ, m.Id()) 181 c.Assert(m.SetProvisioned(inst.Id(), "fake_nonce", md), gc.IsNil) 182 machines[i] = m 183 } 184 return machines 185 } 186 187 func (s *SSHCommonSuite) addUnit(srv *state.Service, m *state.Machine, c *gc.C) { 188 u, err := srv.AddUnit() 189 c.Assert(err, gc.IsNil) 190 err = u.AssignToMachine(m) 191 c.Assert(err, gc.IsNil) 192 // fudge unit.SetPublicAddress 193 id, err := m.InstanceId() 194 c.Assert(err, gc.IsNil) 195 insts, err := s.Conn.Environ.Instances([]instance.Id{id}) 196 c.Assert(err, gc.IsNil) 197 addr, err := insts[0].WaitDNSName() 198 c.Assert(err, gc.IsNil) 199 err = u.SetPublicAddress(addr) 200 c.Assert(err, gc.IsNil) 201 }