github.com/mattyw/juju@v0.0.0-20140610034352-732aecd63861/environs/manual/fakessh_test.go (about)

     1  // Copyright 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package manual_test
     5  
     6  import (
     7  	"fmt"
     8  	"io/ioutil"
     9  	"path/filepath"
    10  	"strings"
    11  
    12  	"github.com/juju/testing"
    13  	gc "launchpad.net/gocheck"
    14  
    15  	"github.com/juju/juju/environs/manual"
    16  )
    17  
    18  // sshscript should only print the result on the first execution,
    19  // to handle the case where it's called multiple times. On
    20  // subsequent executions, it should find the next 'ssh' in $PATH
    21  // and exec that.
    22  var sshscript = `#!/bin/bash --norc
    23  if [ ! -e "$0.run" ]; then
    24      touch "$0.run"
    25      if [ -e "$0.expected-input" ]; then
    26          diff "$0.expected-input" -
    27          exitcode=$?
    28          if [ $exitcode -ne 0 ]; then
    29              echo "ERROR: did not match expected input" >&2
    30              exit $exitcode
    31          fi
    32          else
    33              head >/dev/null
    34          fi
    35      # stdout
    36      %s
    37      # stderr
    38      %s
    39      exit %d
    40  else
    41      export PATH=${PATH#*:}
    42      exec ssh $*
    43  fi`
    44  
    45  // installFakeSSH creates a fake "ssh" command in a new $PATH,
    46  // updates $PATH, and returns a function to reset $PATH to its
    47  // original value when called.
    48  //
    49  // input may be:
    50  //    - nil (ignore input)
    51  //    - a string (match input exactly)
    52  // output may be:
    53  //    - nil (no output)
    54  //    - a string (stdout)
    55  //    - a slice of strings, of length two (stdout, stderr)
    56  func installFakeSSH(c *gc.C, input, output interface{}, rc int) testing.Restorer {
    57  	fakebin := c.MkDir()
    58  	ssh := filepath.Join(fakebin, "ssh")
    59  	switch input := input.(type) {
    60  	case nil:
    61  	case string:
    62  		sshexpectedinput := ssh + ".expected-input"
    63  		err := ioutil.WriteFile(sshexpectedinput, []byte(input), 0644)
    64  		c.Assert(err, gc.IsNil)
    65  	default:
    66  		c.Errorf("input has invalid type: %T", input)
    67  	}
    68  	var stdout, stderr string
    69  	switch output := output.(type) {
    70  	case nil:
    71  	case string:
    72  		stdout = fmt.Sprintf("cat<<EOF\n%s\nEOF", output)
    73  	case []string:
    74  		c.Assert(output, gc.HasLen, 2)
    75  		stdout = fmt.Sprintf("cat<<EOF\n%s\nEOF", output[0])
    76  		stderr = fmt.Sprintf("cat>&2<<EOF\n%s\nEOF", output[1])
    77  	}
    78  	script := fmt.Sprintf(sshscript, stdout, stderr, rc)
    79  	err := ioutil.WriteFile(ssh, []byte(script), 0777)
    80  	c.Assert(err, gc.IsNil)
    81  	return testing.PatchEnvPathPrepend(fakebin)
    82  }
    83  
    84  // installDetectionFakeSSH installs a fake SSH command, which will respond
    85  // to the series/hardware detection script with the specified
    86  // series/arch.
    87  func installDetectionFakeSSH(c *gc.C, series, arch string) testing.Restorer {
    88  	if series == "" {
    89  		series = "precise"
    90  	}
    91  	if arch == "" {
    92  		arch = "amd64"
    93  	}
    94  	detectionoutput := strings.Join([]string{
    95  		series,
    96  		arch,
    97  		"MemTotal: 4096 kB",
    98  		"processor: 0",
    99  	}, "\n")
   100  	return installFakeSSH(c, manual.DetectionScript, detectionoutput, 0)
   101  }
   102  
   103  // fakeSSH wraps the invocation of InstallFakeSSH based on the parameters.
   104  type fakeSSH struct {
   105  	Series string
   106  	Arch   string
   107  
   108  	// Provisioned should be set to true if the fakeSSH script
   109  	// should respond to checkProvisioned with a non-empty result.
   110  	Provisioned bool
   111  
   112  	// exit code for the checkProvisioned script.
   113  	CheckProvisionedExitCode int
   114  
   115  	// exit code for the machine agent provisioning script.
   116  	ProvisionAgentExitCode int
   117  
   118  	// InitUbuntuUser should be set to true if the fakeSSH script
   119  	// should respond to an attempt to initialise the ubuntu user.
   120  	InitUbuntuUser bool
   121  
   122  	// there are conditions other than error in the above
   123  	// that might cause provisioning to not go ahead, such
   124  	// as tools being missing.
   125  	SkipProvisionAgent bool
   126  
   127  	// detection will be skipped if the series/hardware were
   128  	// detected ahead of time. This should always be set to
   129  	// true when testing Bootstrap.
   130  	SkipDetection bool
   131  }
   132  
   133  // install installs fake SSH commands, which will respond to
   134  // manual provisioning/bootstrapping commands with the specified
   135  // output and exit codes.
   136  func (r fakeSSH) install(c *gc.C) testing.Restorer {
   137  	var restore testing.Restorer
   138  	add := func(input, output interface{}, rc int) {
   139  		restore = restore.Add(installFakeSSH(c, input, output, rc))
   140  	}
   141  	if !r.SkipProvisionAgent {
   142  		add(nil, nil, r.ProvisionAgentExitCode)
   143  	}
   144  	if !r.SkipDetection {
   145  		restore.Add(installDetectionFakeSSH(c, r.Series, r.Arch))
   146  	}
   147  	var checkProvisionedOutput interface{}
   148  	if r.Provisioned {
   149  		checkProvisionedOutput = "/etc/init/jujud-machine-0.conf"
   150  	}
   151  	add(manual.CheckProvisionedScript, checkProvisionedOutput, r.CheckProvisionedExitCode)
   152  	if r.InitUbuntuUser {
   153  		add("", nil, 0)
   154  	}
   155  	return restore
   156  }