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  }