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  }