github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/tools/tassert/tassert.go (about)

     1  // Package tassert provides common asserts for tests
     2  /*
     3   * Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package tassert
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"net/http"
    11  	"os"
    12  	"runtime"
    13  	"strings"
    14  	"sync"
    15  	"testing"
    16  	"time"
    17  )
    18  
    19  var (
    20  	fatalities = make(map[string]struct{})
    21  	mu         sync.Mutex
    22  )
    23  
    24  func CheckFatal(tb testing.TB, err error) {
    25  	if err == nil {
    26  		return
    27  	}
    28  	mu.Lock()
    29  	if _, ok := fatalities[tb.Name()]; ok {
    30  		mu.Unlock()
    31  		fmt.Printf("--- %s: duplicate CheckFatal: %v\n", tb.Name(), err) // see #1057
    32  		runtime.Goexit()
    33  	} else {
    34  		fatalities[tb.Name()] = struct{}{}
    35  		mu.Unlock()
    36  		printStack()
    37  		now := fmt.Sprintf("[%s]", time.Now().Format("15:04:05.000000"))
    38  		tb.Fatal(now, err)
    39  	}
    40  }
    41  
    42  func CheckError(tb testing.TB, err error) {
    43  	if err != nil {
    44  		printStack()
    45  		now := fmt.Sprintf("[%s]", time.Now().Format("15:04:05.000000"))
    46  		tb.Error(now, err)
    47  	}
    48  }
    49  
    50  func DoAndCheckResp(tb testing.TB, client *http.Client, req *http.Request, statusCode ...int) {
    51  	resp, err := client.Do(req)
    52  	CheckFatal(tb, err)
    53  	resp.Body.Close()
    54  	for _, code := range statusCode {
    55  		if resp.StatusCode == code {
    56  			return
    57  		}
    58  	}
    59  	Errorf(tb, false, "expected %v status code, got %d", statusCode, resp.StatusCode)
    60  }
    61  
    62  func Fatalf(tb testing.TB, cond bool, msg string, args ...any) {
    63  	if !cond {
    64  		printStack()
    65  		tb.Fatalf(msg, args...)
    66  	}
    67  }
    68  
    69  func Errorf(tb testing.TB, cond bool, msg string, args ...any) {
    70  	if !cond {
    71  		printStack()
    72  		tb.Errorf(msg, args...)
    73  	}
    74  }
    75  
    76  // TODO: Make this a range over `errCh` post closing it ?
    77  func SelectErr(tb testing.TB, errCh chan error, verb string, errIsFatal bool) {
    78  	if num := len(errCh); num > 0 {
    79  		err := <-errCh
    80  		f := tb.Errorf
    81  		if errIsFatal {
    82  			f = tb.Fatalf
    83  		}
    84  		if num > 1 {
    85  			f("Failed to %s %d objects, e.g. error:\n%v", verb, num, err)
    86  		} else {
    87  			f("Failed to %s object: %v", verb, err)
    88  		}
    89  	}
    90  }
    91  
    92  func printStack() {
    93  	var buffer bytes.Buffer
    94  	fmt.Fprintln(os.Stderr, "    tassert.printStack:")
    95  	for i := 1; i < 9; i++ {
    96  		_, file, line, ok := runtime.Caller(i)
    97  		if !ok {
    98  			break
    99  		}
   100  		i := strings.Index(file, "aistore")
   101  		if i < 0 {
   102  			break
   103  		}
   104  		if strings.Contains(file, "tassert") {
   105  			continue
   106  		}
   107  		fmt.Fprintf(&buffer, "\t%s:%d\n", file[i+8:], line)
   108  	}
   109  	os.Stderr.Write(buffer.Bytes())
   110  }