github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/cmds/core/elvish/util/tempdir.go (about)

     1  package util
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"os"
     7  	"path/filepath"
     8  )
     9  
    10  // WithTempDirs creates a requested number of temporary directories and runs a
    11  // function, passing the paths of the temporary directories; the passed paths
    12  // all have their symlinks resolved using filepath.EvalSymlinks. After the
    13  // function returns, it removes the temporary directories. It panics if it
    14  // cannot make a temporary directory, and prints an error message to stderr if
    15  // it cannot remove the temporary directories.
    16  //
    17  // It is useful in tests.
    18  func WithTempDirs(n int, f func([]string)) {
    19  	tmpdirs := make([]string, n)
    20  	for i := range tmpdirs {
    21  		tmpdir, err := ioutil.TempDir("", "elvishtest.")
    22  		if err != nil {
    23  			panic(err)
    24  		}
    25  		tmpdirs[i], err = filepath.EvalSymlinks(tmpdir)
    26  		if err != nil {
    27  			panic(err)
    28  		}
    29  	}
    30  	defer func() {
    31  		for _, tmpdir := range tmpdirs {
    32  			err := os.RemoveAll(tmpdir)
    33  			if err != nil {
    34  				fmt.Fprintln(os.Stderr, "Warning: failed to remove temp dir", tmpdir)
    35  			}
    36  		}
    37  	}()
    38  
    39  	f(tmpdirs)
    40  }
    41  
    42  // WithTempDir is like with WithTempDirs, except that it always creates one
    43  // temporary directory and pass the function a string instead of []string.
    44  func WithTempDir(f func(string)) {
    45  	WithTempDirs(1, func(s []string) {
    46  		f(s[0])
    47  	})
    48  }
    49  
    50  // InTempDir is like WithTempDir, but also cd into the directory before running
    51  // the function, and cd backs after running the function if possible.
    52  //
    53  // It panics if it could not get the working directory or change directory.
    54  //
    55  // It is useful in tests.
    56  func InTempDir(f func(string)) {
    57  	WithTempDir(func(tmpdir string) {
    58  		oldpwd, err := os.Getwd()
    59  		if err != nil {
    60  			panic(err)
    61  		}
    62  
    63  		mustChdir(tmpdir)
    64  		defer mustChdir(oldpwd)
    65  
    66  		f(tmpdir)
    67  	})
    68  }
    69  
    70  func mustChdir(dir string) {
    71  	err := os.Chdir(dir)
    72  	if err != nil {
    73  		panic(err)
    74  	}
    75  }