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