github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/util/clienttest/client_test_suite.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package clienttest
    23  
    24  import (
    25  	"io/ioutil"
    26  	"os"
    27  	"path"
    28  
    29  	flag "github.com/juju/gnuflag"
    30  	"github.com/stretchr/testify/suite"
    31  
    32  	"github.com/dolthub/dolt/go/libraries/utils/osutil"
    33  	"github.com/dolthub/dolt/go/store/d"
    34  	"github.com/dolthub/dolt/go/store/util/exit"
    35  	"github.com/dolthub/dolt/go/store/util/tempfiles"
    36  )
    37  
    38  const DefaultMemTableSize = 8 * (1 << 20) // 8MB
    39  
    40  type ClientTestSuite struct {
    41  	suite.Suite
    42  	TempDir    string
    43  	DBDir      string
    44  	DBDir2     string
    45  	ExitStatus int
    46  	out        *os.File
    47  	err        *os.File
    48  }
    49  
    50  type ExitError struct {
    51  	Code int
    52  }
    53  
    54  func (suite *ClientTestSuite) SetupSuite() {
    55  	td := tempfiles.MovableTempFileProvider.GetTempDir()
    56  	dir, err := ioutil.TempDir(td, "nomstest")
    57  	d.Chk.NoError(err)
    58  	stdOutput, err := ioutil.TempFile(dir, "out")
    59  	d.Chk.NoError(err)
    60  	errOutput, err := ioutil.TempFile(dir, "err")
    61  	d.Chk.NoError(err)
    62  
    63  	suite.TempDir = dir
    64  	suite.DBDir = path.Join(dir, "db")
    65  	suite.DBDir2 = path.Join(suite.TempDir, "db2")
    66  	suite.out = stdOutput
    67  	suite.err = errOutput
    68  	exit.Exit = MockExit
    69  
    70  	os.Mkdir(suite.DBDir, 0777)
    71  	os.Mkdir(suite.DBDir2, 0777)
    72  }
    73  
    74  func (suite *ClientTestSuite) TearDownSuite() {
    75  	suite.out.Close()
    76  	suite.err.Close()
    77  	err := os.RemoveAll(suite.TempDir)
    78  	if !osutil.IsWindows {
    79  		d.Chk.NoError(err)
    80  	}
    81  }
    82  
    83  // MustRun is a wrapper around Run that will panic on Exit or Panic
    84  func (suite *ClientTestSuite) MustRun(m func(), args []string) (stdout string, stderr string) {
    85  	var err interface{}
    86  	if stdout, stderr, err = suite.Run(m, args); err != nil {
    87  		panic(err)
    88  	}
    89  	return
    90  }
    91  
    92  // Run will execute a function passing to it commandline args, and captures stdout,stderr.
    93  // If m()  panics the panic is caught, and returned with recoveredError
    94  // If m() calls exit.Exit() m() will panic and return ExitError with recoveredError
    95  func (suite *ClientTestSuite) Run(m func(), args []string) (stdout string, stderr string, recoveredErr interface{}) {
    96  	origArgs := os.Args
    97  	origOut := os.Stdout
    98  	origErr := os.Stderr
    99  
   100  	os.Args = append([]string{"cmd"}, args...)
   101  	os.Stdout = suite.out
   102  	os.Stderr = suite.err
   103  
   104  	defer func() {
   105  		recoveredErr = recover()
   106  
   107  		// Reset everything right away so that error-checking below goes to terminal.
   108  		os.Args = origArgs
   109  		os.Stdout = origOut
   110  		os.Stderr = origErr
   111  
   112  		_, err := suite.out.Seek(0, 0)
   113  		d.Chk.NoError(err)
   114  		capturedOut, err := ioutil.ReadAll(suite.out)
   115  		d.Chk.NoError(err)
   116  
   117  		_, err = suite.out.Seek(0, 0)
   118  		d.Chk.NoError(err)
   119  		err = suite.out.Truncate(0)
   120  		d.Chk.NoError(err)
   121  
   122  		_, err = suite.err.Seek(0, 0)
   123  		d.Chk.NoError(err)
   124  		capturedErr, err := ioutil.ReadAll(suite.err)
   125  		d.Chk.NoError(err)
   126  
   127  		_, err = suite.err.Seek(0, 0)
   128  		d.Chk.NoError(err)
   129  		err = suite.err.Truncate(0)
   130  		d.Chk.NoError(err)
   131  		stdout, stderr = string(capturedOut), string(capturedErr)
   132  	}()
   133  
   134  	suite.ExitStatus = 0
   135  	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
   136  	m()
   137  	return
   138  }
   139  
   140  // Mock exit.Exit() implementation for use during testing.
   141  func MockExit(status int) {
   142  	panic(ExitError{status})
   143  }