Merge lp:~axwalk/juju-core/sshstorage-tmpdir into lp:~go-bot/juju-core/trunk

[Andrew Wilkins]

environs/sshstorage: revert storage layout

In a lapse of sanity, I changed sshstorage's
layout so that the temporary directory was
contained within the designated storage
directory. This works fine in the context of
sshstorage, but not so well with how sshstorage
is intended to be used.

For bootstrapping the null provider, sshstorage
is used to operate on a specified remote
directory; later, when the machine agent is up,
it takes ownership of that directory and serves
it via environs/localstorage. Thus, both
sshstorage and localstorage must agree on the
layout.

Now, instead of having storage/{tmp,content},
the user can specify a separate temporary
directory. If left blank, storagedir+".tmp" will
be used.

https://codereview.appspot.com/13660047/

environs/sshstorage: revert storage layout

In a lapse of sanity, I changed sshstorage's
layout so that the temporary directory was
contained within the designated storage
directory. This works fine in the context of
sshstorage, but not so well with how sshstorage
is intended to be used.

For bootstrapping the null provider, sshstorage
is used to operate on a specified remote
directory; later, when the machine agent is up,
it takes ownership of that directory and serves
it via environs/localstorage. Thus, both
sshstorage and localstorage must agree on the
layout.

Now, instead of having storage/{tmp,content},
the user can specify a separate temporary
directory. If left blank, storagedir+".tmp" will
be used.

https://codereview.appspot.com/13660047/

[Andrew Wilkins (axwalk)]

Reviewers: mp+185717_code.launchpad.net,

Message:
Please take a look.

Description:
environs/sshstorage: revert storage layout

In a lapse of sanity, I changed sshstorage's
layout so that the temporary directory was
contained within the designated storage
directory. This works fine in the context of
sshstorage, but not so well with how sshstorage
is intended to be used.

For bootstrapping the null provider, sshstorage
is used to operate on a specified remote
directory; later, when the machine agent is up,
it takes ownership of that directory and serves
it via environs/localstorage. Thus, both
sshstorage and localstorage must agree on the
layout.

Now, instead of having storage/{tmp,content},
the user can specify a separate temporary
directory. If left blank, storagedir+".tmp" will
be used.

https://code.launchpad.net/~axwalk/juju-core/sshstorage-tmpdir/+merge/185717

(do not edit description out of merge proposal)

Please review this at https://codereview.appspot.com/13660047/

Affected files (+48, -58 lines):
   A [revision details]
   M environs/sshstorage/storage.go
   M environs/sshstorage/storage_test.go

[Tim Penhey (thumper)]

Mostly comments about test improvements.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go
File environs/sshstorage/storage_test.go (right):

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode80
environs/sshstorage/storage_test.go:80: for i := 0; i < 2; i++ {
It isn't immediately obvious why you are doing this twice.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode81
environs/sshstorage/storage_test.go:81: storage, err :=
NewSSHStorage("example.com", storageDir, "")
Same comment as for the filestorage, the "" here isn't obvious.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode101
environs/sshstorage/storage_test.go:101:
c.Assert(os.Mkdir(filepath.Join(storageDir, "a"), 0755), gc.IsNil)
I have to admit to finding putting the work inside the assert harder to
follow than having them separated out. To me it appears to confuse the
work and the test.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode170
environs/sshstorage/storage_test.go:170: defer storage.Close()
Now that I've seen this same code three times... a function is in order.

func (s *storageSuite) makeStorage(c *gc.C) <storageType> {
   storageDir := c.MkDir()
   storage, err := NewSSHStorage("example.com", storageDir, "")
   c.Assert(err, gc.IsNil)
   s.AddCleanup(func (*gc.C) {
     storage.Close()
   })
   return storage
}

The tests then each get shorter, and logic is in one place.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode177
environs/sshstorage/storage_test.go:177:
c.Assert(ioutil.WriteFile(filepath.Join(storageDir, "a", "b1"), nil, 0),
gc.IsNil)
This is also screaming out for a helper method.

You want the method to describe what is happening.

https://codereview.appspot.com/13660047/

[Andrew Wilkins (axwalk)]

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go
File environs/sshstorage/storage_test.go (right):

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode80
environs/sshstorage/storage_test.go:80: for i := 0; i < 2; i++ {
On 2013/09/18 04:39:41, thumper wrote:
> It isn't immediately obvious why you are doing this twice.

Added a comment.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode81
environs/sshstorage/storage_test.go:81: storage, err :=
NewSSHStorage("example.com", storageDir, "")
On 2013/09/18 04:39:41, thumper wrote:
> Same comment as for the filestorage, the "" here isn't obvious.

Done. Also updated Put to create/remove remove the default temporary
directory, though it's less of an issue here.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode101
environs/sshstorage/storage_test.go:101:
c.Assert(os.Mkdir(filepath.Join(storageDir, "a"), 0755), gc.IsNil)
On 2013/09/18 04:39:41, thumper wrote:
> I have to admit to finding putting the work inside the assert harder
to follow
> than having them separated out. To me it appears to confuse the work
and the
> test.

I've separated the lengthier ones out to make it obvious what's the
code-under-test and what's the assertion; I've left many of the shorter
assertions alone, e.g. c.Assert(storage.Remove("a"), gc.NotNil).

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode170
environs/sshstorage/storage_test.go:170: defer storage.Close()
On 2013/09/18 04:39:41, thumper wrote:
> Now that I've seen this same code three times... a function is in
order.

> func (s *storageSuite) makeStorage(c *gc.C) <storageType> {
> storageDir := c.MkDir()
> storage, err := NewSSHStorage("example.com", storageDir, "")
> c.Assert(err, gc.IsNil)
> s.AddCleanup(func (*gc.C) {
> storage.Close()
> })
> return storage
> }

> The tests then each get shorter, and logic is in one place.

Done.

https://codereview.appspot.com/13660047/diff/1/environs/sshstorage/storage_test.go#newcode177
environs/sshstorage/storage_test.go:177:
c.Assert(ioutil.WriteFile(filepath.Join(storageDir, "a", "b1"), nil, 0),
gc.IsNil)
On 2013/09/18 04:39:41, thumper wrote:
> This is also screaming out for a helper method.

> You want the method to describe what is happening.

Yep, done. Added a "createFiles" test helper.

https://codereview.appspot.com/13660047/

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/13660047/

[Tim Penhey (thumper)]

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go
File environs/sshstorage/storage_test.go (right):

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go#newcode4
environs/sshstorage/storage_test.go:4: package sshstorage
Normally the tests would be in sshstorage_test package.

Is there any reason why this is a whitebox test?

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go#newcode74
environs/sshstorage/storage_test.go:74: func (s *storageSuite)
makeStorage(c *gc.C) (storage *SSHStorage, storageDir string) {
Why name the return parameters here? You aren't really gaining much. Is
it just to provide more assistance to the reader?

https://codereview.appspot.com/13660047/

[Tim Penhey (thumper)]

review: Needs Information

[Andrew Wilkins (axwalk)]

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go
File environs/sshstorage/storage_test.go (right):

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go#newcode4
environs/sshstorage/storage_test.go:4: package sshstorage
On 2013/09/18 21:30:56, thumper wrote:
> Normally the tests would be in sshstorage_test package.

> Is there any reason why this is a whitebox test?

Yes, so the tests can mock out "sshCommand". See the call to jc.Set in
SetUpSuite.

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go#newcode74
environs/sshstorage/storage_test.go:74: func (s *storageSuite)
makeStorage(c *gc.C) (storage *SSHStorage, storageDir string) {
On 2013/09/18 21:30:56, thumper wrote:
> Why name the return parameters here? You aren't really gaining much.
Is it just
> to provide more assistance to the reader?

Yes, that was my rationale; I don't want to have to look at the
implementation to know what the result is.

It could be a doc comment, but in this case a self-documenting named
result is an easy win.

https://codereview.appspot.com/13660047/

[Tim Penhey (thumper)]

for tests that are whitebox tests, we have had an informal convention of
naming them _whitebox_test.go instead of _test.go

LGTM

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go
File environs/sshstorage/storage_test.go (right):

https://codereview.appspot.com/13660047/diff/7001/environs/sshstorage/storage_test.go#newcode4
environs/sshstorage/storage_test.go:4: package sshstorage
On 2013/09/19 02:36:19, axw1 wrote:
> On 2013/09/18 21:30:56, thumper wrote:
> > Normally the tests would be in sshstorage_test package.
> >
> > Is there any reason why this is a whitebox test?

> Yes, so the tests can mock out "sshCommand". See the call to jc.Set in
> SetUpSuite.

jc.Set has moved, and is about to move again with my branch :-)

   testing.PatchValue
until it becomes
   testbase.PatchValue

https://codereview.appspot.com/13660047/

[Go Bot (go-bot)]

The attempt to merge lp:~axwalk/juju-core/sshstorage-tmpdir into lp:juju-core failed. Below is the output from the failed tests.

ok launchpad.net/juju-core/agent 0.743s
ok launchpad.net/juju-core/agent/tools 0.275s
ok launchpad.net/juju-core/bzr 6.637s
ok launchpad.net/juju-core/cert 2.373s
ok launchpad.net/juju-core/charm 0.541s
? launchpad.net/juju-core/charm/hooks [no test files]
ok launchpad.net/juju-core/cloudinit 0.021s
ok launchpad.net/juju-core/cmd 0.244s
? launchpad.net/juju-core/cmd/builddb [no test files]
? launchpad.net/juju-core/cmd/charmd [no test files]
? launchpad.net/juju-core/cmd/charmload [no test files]
ok launchpad.net/juju-core/cmd/juju 154.458s
ok launchpad.net/juju-core/cmd/jujud 39.250s
ok launchpad.net/juju-core/cmd/plugins/juju-metadata 1.840s
ok launchpad.net/juju-core/constraints 0.028s
ok launchpad.net/juju-core/container/lxc 0.351s
? launchpad.net/juju-core/container/lxc/mock [no test files]
ok launchpad.net/juju-core/downloader 5.288s
ok launchpad.net/juju-core/environs 1.844s
ok launchpad.net/juju-core/environs/bootstrap 4.872s
ok launchpad.net/juju-core/environs/cloudinit 0.426s
ok launchpad.net/juju-core/environs/config 0.799s
ok launchpad.net/juju-core/environs/configstore 0.063s
ok launchpad.net/juju-core/environs/filestorage 0.030s
ok launchpad.net/juju-core/environs/httpstorage 0.048s
ok launchpad.net/juju-core/environs/imagemetadata 0.489s
ok launchpad.net/juju-core/environs/instances 0.052s
ok launchpad.net/juju-core/environs/jujutest 0.214s
ok launchpad.net/juju-core/environs/manual 1.115s
ok launchpad.net/juju-core/environs/simplestreams 0.265s
? launchpad.net/juju-core/environs/simplestreams/testing [no test files]
FAIL launchpad.net/juju-core/environs/sshstorage [build failed]
ok launchpad.net/juju-core/environs/storage 0.011s
ok launchpad.net/juju-core/environs/sync 26.073s
ok launchpad.net/juju-core/environs/testing 0.011s
ok launchpad.net/juju-core/environs/tools 10.424s
? launchpad.net/juju-core/environs/tools/testing [no test files]
? launchpad.net/juju-core/errors [no test files]
ok launchpad.net/juju-core/instance 0.022s
? launchpad.net/juju-core/instance/testing [no test files]
ok launchpad.net/juju-core/juju 20.195s
ok launchpad.net/juju-core/juju/osenv 0.018s
? launchpad.net/juju-core/juju/testing [no test files]
ok launchpad.net/juju-core/log 0.015s
ok launchpad.net/juju-core/log/syslog 0.022s
ok launchpad.net/juju-core/names 0.024s
ok launchpad.net/juju-core/provider 0.328s
? launchpad.net/juju-core/provider/all [no test files]
ok launchpad.net/juju-core/provider/azure 6.439s
ok launchpad.net/juju-core/provider/dummy 20.225s
ok launchpad.net/juju-core/provider/ec2 5.228s
ok launchpad.net/juju-core/provider/ec2/httpstorage 0.050s
ok launchpad.net/juju-core/provider/local 1.779s
? launchpad.net/juju-core/provider/local/storage [no test files]
ok launchpad.net/juju-core/provider/maas 3.733s
ok launchpad.net/juju-core/provider/openstack 8.343s
ok launchpad.net/juju-core/rpc 0.080s
ok launchpad.net/juju-core/rpc/jsoncodec 0.038s
ok launchpad.net/juju-core/schema 0....

=== modified file 'environs/sshstorage/storage.go'
--- environs/sshstorage/storage.go	2013-09-18 22:54:32 +0000
+++ environs/sshstorage/storage.go	2013-09-20 04:07:37 +0000
@@ -21,18 +21,6 @@
 	"launchpad.net/juju-core/utils"
 )
 
-const (
-	// tmpdir is the name of the subdirectory
-	// inside the remote storage directory where
-	// temporary files are created.
-	tmpdir = "tmp"
-
-	// contentdir is the name of the subdirectory
-	// inside the remote storage directory where
-	// files are stored.
-	contentdir = "content"
-)
-
 // SSHStorage implements storage.Storage.
 //
 // The storage is created under sudo, and ownership given over to the
@@ -42,6 +30,7 @@
 type SSHStorage struct {
 	host       string
 	remotepath string
+	tmpdir     string
 
 	cmd     *exec.Cmd
 	stdin   io.WriteCloser
@@ -65,12 +54,24 @@
 	flockExclusive flockmode = "-x"
 )
 
+// UseDefaultTmpDir may be passed into NewSSHStorage
+// for the tmpdir argument, to signify that the default
+// value should be used. See NewSSHStorage for more.
+const UseDefaultTmpDir = ""
+
 // NewSSHStorage creates a new SSHStorage, connected to the
 // specified host, managing state under the specified remote path.
-func NewSSHStorage(host string, remotepath string) (*SSHStorage, error) {
-	contentdir := path.Join(remotepath, contentdir)
-	tmpdir := path.Join(remotepath, tmpdir)
-	script := fmt.Sprintf("install -d -g $SUDO_GID -o $SUDO_UID %s %s", contentdir, tmpdir)
+//
+// A temporary directory may be specified, in which case it should
+// be a directory on the same filesystem as the storage directory
+// to ensure atomic writes. If left unspecified, tmpdir will be
+// assigned a value of storagedir+".tmp".
+//
+// If tmpdir == UseDefaultTmpDir, it will be created when Put is invoked,
+// and will be removed afterwards. If tmpdir != UseDefaultTmpDir, it must
+// already exist, and will never be removed.
+func NewSSHStorage(host, storagedir, tmpdir string) (*SSHStorage, error) {
+	script := fmt.Sprintf("install -d -g $SUDO_GID -o $SUDO_UID %s", utils.ShQuote(storagedir))
 	cmd := sshCommand(host, true, fmt.Sprintf("sudo bash -c '%s'", script))
 	cmd.Stdin = os.Stdin
 	if out, err := cmd.CombinedOutput(); err != nil {
@@ -93,7 +94,8 @@
 	}
 	stor := &SSHStorage{
 		host:       host,
-		remotepath: remotepath,
+		remotepath: storagedir,
+		tmpdir:     tmpdir,
 		cmd:        cmd,
 		stdin:      stdin,
 		stdout:     stdout,
@@ -101,13 +103,8 @@
 	}
 	cmd.Start()
 
-	// Verify we have write permissions, and set the temporary directory.
-	_, err = stor.runf(
-		flockExclusive,
-		"touch %s && export TMPDIR=%s",
-		utils.ShQuote(remotepath),
-		utils.ShQuote(tmpdir),
-	)
+	// Verify we have write permissions.
+	_, err = stor.runf(flockExclusive, "touch %s", utils.ShQuote(storagedir))
 	if err != nil {
 		stdin.Close()
 		stdout.Close()
@@ -164,9 +161,8 @@
 
 // path returns a remote absolute path for a storage object name.
 func (s *SSHStorage) path(name string) (string, error) {
-	contentdir := path.Join(s.remotepath, contentdir)
-	remotepath := path.Clean(path.Join(contentdir, name))
-	if !strings.HasPrefix(remotepath, contentdir) {
+	remotepath := path.Clean(path.Join(s.remotepath, name))
+	if !strings.HasPrefix(remotepath, s.remotepath) {
 		return "", fmt.Errorf("%q escapes storage directory", name)
 	}
 	return remotepath, nil
@@ -209,10 +205,9 @@
 		return nil, nil
 	}
 	var names []string
-	contentdir := path.Join(s.remotepath, contentdir)
 	for _, name := range strings.Split(out, "\n") {
 		if strings.HasPrefix(name[len(dir):], prefix) {
-			names = append(names, name[len(contentdir)+1:])
+			names = append(names, name[len(s.remotepath)+1:])
 		}
 	}
 	sort.Strings(names)
@@ -250,9 +245,30 @@
 	}
 	encoded := base64.StdEncoding.EncodeToString(buf)
 	path = utils.ShQuote(path)
+
+	tmpdir := s.tmpdir
+	if tmpdir == UseDefaultTmpDir {
+		tmpdir = s.remotepath + ".tmp"
+	}
+	tmpdir = utils.ShQuote(tmpdir)
+
 	// Write to a temporary file ($TMPFILE), then mv atomically.
 	command := fmt.Sprintf("mkdir -p `dirname %s` && base64 -d > $TMPFILE", path)
-	command = fmt.Sprintf("TMPFILE=`mktemp` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)", command, path)
+	command = fmt.Sprintf(
+		"export TMPDIR=%s && TMPFILE=`mktemp` && ((%s && mv $TMPFILE %s) || rm -f $TMPFILE)",
+		tmpdir, command, path,
+	)
+
+	// If UseDefaultTmpDir is passed, then create the
+	// temporary directory, and remove it afterwards.
+	// Otherwise, the temporary directory is expected
+	// to already exist, and is never removed.
+	if s.tmpdir == UseDefaultTmpDir {
+		installTmpdir := fmt.Sprintf("install -d -g $SUDO_GID -o $SUDO_UID %s", tmpdir)
+		removeTmpdir := fmt.Sprintf("rm -fr %s", tmpdir)
+		command = fmt.Sprintf("%s && (%s); rc=$?; %s; exit $rc", installTmpdir, command, removeTmpdir)
+	}
+
 	command = fmt.Sprintf("(%s) << EOF\n%s\nEOF", command, encoded)
 	_, err = s.runf(flockExclusive, command+"\n")
 	return err
@@ -271,7 +287,6 @@
 
 // RemoveAll implements storage.StorageWriter.RemoveAll
 func (s *SSHStorage) RemoveAll() error {
-	contentdir := path.Join(s.remotepath, contentdir)
-	_, err := s.runf(flockExclusive, "rm -fr %s/*", utils.ShQuote(contentdir))
+	_, err := s.runf(flockExclusive, "rm -fr %s/*", utils.ShQuote(s.remotepath))
 	return err
 }

=== modified file 'environs/sshstorage/storage_test.go'
--- environs/sshstorage/storage_test.go	2013-09-20 02:53:59 +0000
+++ environs/sshstorage/storage_test.go	2013-09-20 04:07:37 +0000
@@ -27,7 +27,6 @@
 
 type storageSuite struct {
 	testbase.LoggingSuite
-	restoreEnv func()
 }
 
 var _ = gc.Suite(&storageSuite{})
@@ -36,8 +35,6 @@
 	gc.TestingT(t)
 }
 
-var sshCommandOrig = sshCommand
-
 func sshCommandTesting(host string, tty bool, command string) *exec.Cmd {
 	cmd := exec.Command("bash", "-c", command)
 	uid := fmt.Sprint(os.Getuid())
@@ -59,50 +56,70 @@
 	c.Assert(err, gc.IsNil)
 
 	bin := c.MkDir()
-	s.restoreEnv = testbase.PatchEnvironment("PATH", bin+":"+os.Getenv("PATH"))
+	restoreEnv := testbase.PatchEnvironment("PATH", bin+":"+os.Getenv("PATH"))
+	s.AddSuiteCleanup(func(*gc.C) { restoreEnv() })
 
 	// Create a "sudo" command which just executes its args.
-	c.Assert(os.Symlink("/usr/bin/env", filepath.Join(bin, "sudo")), gc.IsNil)
-	sshCommand = sshCommandTesting
+	err = os.Symlink("/usr/bin/env", filepath.Join(bin, "sudo"))
+	c.Assert(err, gc.IsNil)
+	restoreSshCommand := testbase.PatchValue(&sshCommand, sshCommandTesting)
+	s.AddSuiteCleanup(func(*gc.C) { restoreSshCommand() })
 
 	// Create a new "flock" which calls the original, but in non-blocking mode.
 	data := []byte(fmt.Sprintf("#!/bin/sh\nexec %s --nonblock \"$@\"", flockBin))
-	c.Assert(ioutil.WriteFile(filepath.Join(bin, "flock"), data, 0755), gc.IsNil)
-}
-
-func (s *storageSuite) TearDownSuite(c *gc.C) {
-	sshCommand = sshCommandOrig
-	s.restoreEnv()
-	s.LoggingSuite.TearDownSuite(c)
+	err = ioutil.WriteFile(filepath.Join(bin, "flock"), data, 0755)
+	c.Assert(err, gc.IsNil)
+}
+
+func (s *storageSuite) makeStorage(c *gc.C) (storage *SSHStorage, storageDir string) {
+	storageDir = c.MkDir()
+	storage, err := NewSSHStorage("example.com", storageDir, "")
+	c.Assert(err, gc.IsNil)
+	c.Assert(storage, gc.NotNil)
+	s.AddCleanup(func(*gc.C) { storage.Close() })
+	return storage, storageDir
+}
+
+// createFiles creates empty files in the storage directory
+// with the given storage names.
+func createFiles(c *gc.C, storageDir string, names ...string) {
+	for _, name := range names {
+		path := filepath.Join(storageDir, filepath.FromSlash(name))
+		dir := filepath.Dir(path)
+		if err := os.MkdirAll(dir, 0755); err != nil {
+			c.Assert(err, jc.Satisfies, os.IsExist)
+		}
+		err := ioutil.WriteFile(path, nil, 0644)
+		c.Assert(err, gc.IsNil)
+	}
 }
 
 func (s *storageSuite) TestNewSSHStorage(c *gc.C) {
 	storageDir := c.MkDir()
+	// Run this block twice to ensure NewSSHStorage can reuse
+	// an existing storage location.
 	for i := 0; i < 2; i++ {
-		stor, err := NewSSHStorage("example.com", storageDir)
+		stor, err := NewSSHStorage("example.com", storageDir, UseDefaultTmpDir)
 		c.Assert(err, gc.IsNil)
 		c.Assert(stor, gc.NotNil)
 		c.Assert(stor.Close(), gc.IsNil)
 	}
-	c.Assert(os.RemoveAll(storageDir), gc.IsNil)
+	err := os.RemoveAll(storageDir)
+	c.Assert(err, gc.IsNil)
 
 	// You must have permissions to create the directory.
 	storageDir = c.MkDir()
-	c.Assert(os.Chmod(storageDir, 0555), gc.IsNil)
-	_, err := NewSSHStorage("example.com", filepath.Join(storageDir))
+	err = os.Chmod(storageDir, 0555)
+	c.Assert(err, gc.IsNil)
+	_, err = NewSSHStorage("example.com", filepath.Join(storageDir, "subdir"), UseDefaultTmpDir)
 	c.Assert(err, gc.ErrorMatches, "(.|\n)*cannot change owner and permissions of(.|\n)*")
 }
 
 func (s *storageSuite) TestPathValidity(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
-	defer stor.Close()
-
-	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
-	f, err := os.Create(filepath.Join(storageDir, contentdir, "a", "b"))
-	c.Assert(err, gc.IsNil)
-	c.Assert(f.Close(), gc.IsNil)
+	stor, storageDir := s.makeStorage(c)
+	err := os.Mkdir(filepath.Join(storageDir, "a"), 0755)
+	c.Assert(err, gc.IsNil)
+	createFiles(c, storageDir, "a/b")
 
 	for _, prefix := range []string{"..", "a/../.."} {
 		c.Logf("prefix: %q", prefix)
@@ -122,14 +139,12 @@
 }
 
 func (s *storageSuite) TestGet(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
+	stor, storageDir := s.makeStorage(c)
+	data := []byte("abc\000def")
+	err := os.Mkdir(filepath.Join(storageDir, "a"), 0755)
 	c.Assert(err, gc.IsNil)
-	defer stor.Close()
-	data := []byte("abc\000def")
-	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
 	for _, name := range []string{"b", filepath.Join("a", "b")} {
-		err = ioutil.WriteFile(filepath.Join(storageDir, contentdir, name), data, 0644)
+		err = ioutil.WriteFile(filepath.Join(storageDir, name), data, 0644)
 		c.Assert(err, gc.IsNil)
 		r, err := storage.Get(stor, name)
 		c.Assert(err, gc.IsNil)
@@ -137,21 +152,17 @@
 		c.Assert(err, gc.IsNil)
 		c.Assert(out, gc.DeepEquals, data)
 	}
-
 	_, err = storage.Get(stor, "notthere")
 	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)
 }
 
 func (s *storageSuite) TestPut(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
-	defer stor.Close()
+	stor, storageDir := s.makeStorage(c)
 	data := []byte("abc\000def")
 	for _, name := range []string{"b", filepath.Join("a", "b")} {
-		err = stor.Put(name, bytes.NewBuffer(data), int64(len(data)))
+		err := stor.Put(name, bytes.NewBuffer(data), int64(len(data)))
 		c.Assert(err, gc.IsNil)
-		out, err := ioutil.ReadFile(filepath.Join(storageDir, contentdir, name))
+		out, err := ioutil.ReadFile(filepath.Join(storageDir, name))
 		c.Assert(err, gc.IsNil)
 		c.Assert(out, gc.DeepEquals, data)
 	}
@@ -165,20 +176,15 @@
 }
 
 func (s *storageSuite) TestList(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
-	defer stor.Close()
+	stor, storageDir := s.makeStorage(c)
 	s.assertList(c, stor, "", nil)
 
 	// Directories don't show up in List.
-	contentDir := filepath.Join(storageDir, contentdir)
-	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
+	err := os.Mkdir(filepath.Join(storageDir, "a"), 0755)
+	c.Assert(err, gc.IsNil)
 	s.assertList(c, stor, "", nil)
 	s.assertList(c, stor, "a", nil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "b"), nil, 0), gc.IsNil)
+	createFiles(c, storageDir, "a/b1", "a/b2", "b")
 	s.assertList(c, stor, "", []string{"a/b1", "a/b2", "b"})
 	s.assertList(c, stor, "a", []string{"a/b1", "a/b2"})
 	s.assertList(c, stor, "a/b", []string{"a/b1", "a/b2"})
@@ -189,15 +195,10 @@
 }
 
 func (s *storageSuite) TestRemove(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
+	stor, storageDir := s.makeStorage(c)
+	err := os.Mkdir(filepath.Join(storageDir, "a"), 0755)
 	c.Assert(err, gc.IsNil)
-	defer stor.Close()
-
-	contentDir := filepath.Join(storageDir, contentdir)
-	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
+	createFiles(c, storageDir, "a/b1", "a/b2")
 	c.Assert(stor.Remove("a"), gc.ErrorMatches, "rm: cannot remove.*Is a directory")
 	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
 	c.Assert(stor.Remove("a/b"), gc.IsNil) // doesn't exist; not an error
@@ -209,39 +210,28 @@
 }
 
 func (s *storageSuite) TestRemoveAll(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
+	stor, storageDir := s.makeStorage(c)
+	err := os.Mkdir(filepath.Join(storageDir, "a"), 0755)
 	c.Assert(err, gc.IsNil)
-	defer stor.Close()
-
-	contentDir := filepath.Join(storageDir, contentdir)
-	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
-	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
+	createFiles(c, storageDir, "a/b1", "a/b2")
 	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
 	c.Assert(stor.RemoveAll(), gc.IsNil)
 	s.assertList(c, stor, "", nil)
 
 	// RemoveAll does not remove the base storage directory.
-	_, err = os.Stat(contentDir)
+	_, err = os.Stat(storageDir)
 	c.Assert(err, gc.IsNil)
 }
 
 func (s *storageSuite) TestURL(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
-	defer stor.Close()
+	stor, storageDir := s.makeStorage(c)
 	url, err := stor.URL("a/b")
 	c.Assert(err, gc.IsNil)
-	c.Assert(url, gc.Equals, "sftp://example.com/"+path.Join(storageDir, contentdir, "a/b"))
+	c.Assert(url, gc.Equals, "sftp://example.com/"+path.Join(storageDir, "a/b"))
 }
 
-func (s *storageSuite) TestConsistencyStrategy(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
-	defer stor.Close()
+func (s *storageSuite) TestDefaultConsistencyStrategy(c *gc.C) {
+	stor, _ := s.makeStorage(c)
 	c.Assert(stor.DefaultConsistencyStrategy(), gc.Equals, utils.AttemptStrategy{})
 }
 
@@ -271,22 +261,19 @@
 	//
 	// flock exits with exit code 1 if it can't acquire the
 	// lock immediately in non-blocking mode (which the tests force).
-	_, err := NewSSHStorage("example.com", storageDir)
+	_, err := NewSSHStorage("example.com", storageDir, UseDefaultTmpDir)
 	c.Assert(err, gc.ErrorMatches, "exit code 1")
 }
 
 func (s *storageSuite) TestWithSharedLocks(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
+	stor, storageDir := s.makeStorage(c)
 
 	// Get and List should be able to proceed with a shared lock.
 	// All other methods should fail.
-	data := []byte("abc\000def")
-	c.Assert(ioutil.WriteFile(filepath.Join(storageDir, contentdir, "a"), data, 0644), gc.IsNil)
+	createFiles(c, storageDir, "a")
 
 	s.flock(c, flockShared, storageDir)
-	_, err = storage.Get(stor, "a")
+	_, err := storage.Get(stor, "a")
 	c.Assert(err, gc.IsNil)
 	_, err = storage.List(stor, "")
 	c.Assert(err, gc.IsNil)
@@ -296,13 +283,11 @@
 }
 
 func (s *storageSuite) TestWithExclusiveLocks(c *gc.C) {
-	storageDir := c.MkDir()
-	stor, err := NewSSHStorage("example.com", storageDir)
-	c.Assert(err, gc.IsNil)
+	stor, storageDir := s.makeStorage(c)
 	// None of the methods (apart from URL) should be able to do anything
 	// while an exclusive lock is held.
 	s.flock(c, flockExclusive, storageDir)
-	_, err = stor.URL("a")
+	_, err := stor.URL("a")
 	c.Assert(err, gc.IsNil)
 	c.Assert(stor.Put("a", bytes.NewBuffer(nil), 0), gc.NotNil)
 	c.Assert(stor.Remove("a"), gc.NotNil)
@@ -312,3 +297,39 @@
 	_, err = storage.List(stor, "")
 	c.Assert(err, gc.NotNil)
 }
+
+func (s *storageSuite) TestPutTmpDir(c *gc.C) {
+	stor, storageDir := s.makeStorage(c)
+
+	// Put should create and clean up the temporary directory if
+	// tmpdir==UseDefaultTmpDir.
+	err := stor.Put("test-write", bytes.NewReader(nil), 0)
+	c.Assert(err, gc.IsNil)
+	_, err = os.Stat(storageDir + ".tmp")
+	c.Assert(err, jc.Satisfies, os.IsNotExist)
+
+	// To deal with recovering from hard failure, UseDefaultTmpDir
+	// doesn't care if the temporary directory already exists. It
+	// still removes it, though.
+	err = os.Mkdir(storageDir+".tmp", 0755)
+	c.Assert(err, gc.IsNil)
+	err = stor.Put("test-write", bytes.NewReader(nil), 0)
+	c.Assert(err, gc.IsNil)
+	_, err = os.Stat(storageDir + ".tmp")
+	c.Assert(err, jc.Satisfies, os.IsNotExist)
+
+	// If we explicitly set the temporary directory, it must already exist.
+	stor.Close()
+	stor, err = NewSSHStorage("example.com", storageDir, storageDir+".tmp")
+	defer stor.Close()
+	c.Assert(err, gc.IsNil)
+	err = stor.Put("test-write", bytes.NewReader(nil), 0)
+	c.Assert(err, gc.ErrorMatches, "mktemp: failed to create file.*No such file or directory")
+	err = os.Mkdir(storageDir+".tmp", 0755)
+	c.Assert(err, gc.IsNil)
+	err = stor.Put("test-write", bytes.NewReader(nil), 0)
+	c.Assert(err, gc.IsNil)
+	// Temporary directory should not have been moved.
+	_, err = os.Stat(storageDir + ".tmp")
+	c.Assert(err, gc.IsNil)
+}