github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/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"
    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/file"
    33  	"github.com/dolthub/dolt/go/libraries/utils/osutil"
    34  	"github.com/dolthub/dolt/go/store/d"
    35  	"github.com/dolthub/dolt/go/store/util/exit"
    36  	"github.com/dolthub/dolt/go/store/util/tempfiles"
    37  )
    38  
    39  const DefaultMemTableSize = 8 * (1 << 20) // 8MB
    40  
    41  type ClientTestSuite struct {
    42  	suite.Suite
    43  	TempDir    string
    44  	DBDir      string
    45  	DBDir2     string
    46  	ExitStatus int
    47  	out        *os.File
    48  	err        *os.File
    49  }
    50  
    51  type ExitError struct {
    52  	Code int
    53  }
    54  
    55  func (suite *ClientTestSuite) SetupSuite() {
    56  	td := tempfiles.MovableTempFileProvider.GetTempDir()
    57  	dir, err := os.MkdirTemp(td, "nomstest")
    58  	d.Chk.NoError(err)
    59  	stdOutput, err := os.CreateTemp(dir, "out")
    60  	d.Chk.NoError(err)
    61  	errOutput, err := os.CreateTemp(dir, "err")
    62  	d.Chk.NoError(err)
    63  
    64  	suite.TempDir = dir
    65  	suite.DBDir = path.Join(dir, "db")
    66  	suite.DBDir2 = path.Join(suite.TempDir, "db2")
    67  	suite.out = stdOutput
    68  	suite.err = errOutput
    69  	exit.Exit = MockExit
    70  
    71  	os.Mkdir(suite.DBDir, 0777)
    72  	os.Mkdir(suite.DBDir2, 0777)
    73  }
    74  
    75  func (suite *ClientTestSuite) TearDownSuite() {
    76  	suite.out.Close()
    77  	suite.err.Close()
    78  	err := file.RemoveAll(suite.TempDir)
    79  	if !osutil.IsWindows {
    80  		d.Chk.NoError(err)
    81  	}
    82  }
    83  
    84  // MustRun is a wrapper around Run that will panic on Exit or Panic
    85  func (suite *ClientTestSuite) MustRun(m func(), args []string) (stdout string, stderr string) {
    86  	var err interface{}
    87  	if stdout, stderr, err = suite.Run(m, args); err != nil {
    88  		panic(err)
    89  	}
    90  	return
    91  }
    92  
    93  // Run will execute a function passing to it commandline args, and captures stdout,stderr.
    94  // If m()  panics the panic is caught, and returned with recoveredError
    95  // If m() calls exit.Exit() m() will panic and return ExitError with recoveredError
    96  func (suite *ClientTestSuite) Run(m func(), args []string) (stdout string, stderr string, recoveredErr interface{}) {
    97  	origArgs := os.Args
    98  	origOut := os.Stdout
    99  	origErr := os.Stderr
   100  
   101  	os.Args = append([]string{"cmd"}, args...)
   102  	os.Stdout = suite.out
   103  	os.Stderr = suite.err
   104  
   105  	defer func() {
   106  		recoveredErr = recover()
   107  
   108  		// Reset everything right away so that error-checking below goes to terminal.
   109  		os.Args = origArgs
   110  		os.Stdout = origOut
   111  		os.Stderr = origErr
   112  
   113  		_, err := suite.out.Seek(0, 0)
   114  		d.Chk.NoError(err)
   115  		capturedOut, err := io.ReadAll(suite.out)
   116  		d.Chk.NoError(err)
   117  
   118  		_, err = suite.out.Seek(0, 0)
   119  		d.Chk.NoError(err)
   120  		err = suite.out.Truncate(0)
   121  		d.Chk.NoError(err)
   122  
   123  		_, err = suite.err.Seek(0, 0)
   124  		d.Chk.NoError(err)
   125  		capturedErr, err := io.ReadAll(suite.err)
   126  		d.Chk.NoError(err)
   127  
   128  		_, err = suite.err.Seek(0, 0)
   129  		d.Chk.NoError(err)
   130  		err = suite.err.Truncate(0)
   131  		d.Chk.NoError(err)
   132  		stdout, stderr = string(capturedOut), string(capturedErr)
   133  	}()
   134  
   135  	suite.ExitStatus = 0
   136  	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
   137  	m()
   138  	return
   139  }
   140  
   141  // Mock exit.Exit() implementation for use during testing.
   142  func MockExit(status int) {
   143  	panic(ExitError{status})
   144  }