github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ais/test/main_test.go (about)

     1  // Package integration_test.
     2  /*
     3   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package integration_test
     6  
     7  import (
     8  	"flag"
     9  	"fmt"
    10  	"net/http"
    11  	"os"
    12  	"strconv"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/NVIDIA/aistore/api"
    17  	"github.com/NVIDIA/aistore/api/apc"
    18  	"github.com/NVIDIA/aistore/api/env"
    19  	"github.com/NVIDIA/aistore/cmn"
    20  	"github.com/NVIDIA/aistore/tools"
    21  	"github.com/NVIDIA/aistore/tools/tassert"
    22  	"github.com/NVIDIA/aistore/tools/tlog"
    23  	"github.com/NVIDIA/aistore/tools/trand"
    24  )
    25  
    26  var (
    27  	proxyURL             string
    28  	baseParams           api.BaseParams
    29  	initialClusterConfig *cmn.ClusterConfig
    30  )
    31  
    32  func setBucket() (bck cmn.Bck, err error) {
    33  	bucket := os.Getenv("BUCKET")
    34  	if bucket == "" {
    35  		bucket = apc.AIS + apc.BckProviderSeparator + trand.String(7)
    36  	}
    37  	bck, _, err = cmn.ParseBckObjectURI(bucket, cmn.ParseURIOpts{})
    38  	if err != nil {
    39  		return bck, fmt.Errorf("failed to parse 'BUCKET' env variable, err: %v", err)
    40  	} else if err := bck.Validate(); err != nil {
    41  		return bck, fmt.Errorf("failed to validate 'BUCKET' env variable, err: %v", err)
    42  	}
    43  	tlog.Logf("Using bucket %s\n", bck)
    44  	return bck, nil
    45  }
    46  
    47  func waitForCluster() (primaryURL string, err error) {
    48  	const (
    49  		retryCount = 30
    50  		sleep      = time.Second
    51  	)
    52  	var (
    53  		proxyCnt, targetCnt, retry int
    54  	)
    55  	pc := os.Getenv(env.AIS.NumProxy)
    56  	tc := os.Getenv(env.AIS.NumTarget)
    57  	if pc != "" || tc != "" {
    58  		proxyCnt, err = strconv.Atoi(pc)
    59  		if err != nil {
    60  			err = fmt.Errorf("error EnvVars: %s. err: %v", env.AIS.NumProxy, err)
    61  			return
    62  		}
    63  		targetCnt, err = strconv.Atoi(tc)
    64  		if err != nil {
    65  			err = fmt.Errorf("error EnvVars: %s. err: %v", env.AIS.NumTarget, err)
    66  			return
    67  		}
    68  	}
    69  	_, err = tools.WaitForClusterState(tools.GetPrimaryURL(), "cluster startup", -1, proxyCnt, targetCnt)
    70  	if err != nil {
    71  		err = fmt.Errorf("error waiting for cluster startup, err: %v", err)
    72  		return
    73  	}
    74  	tlog.Logf("Pinging primary for readiness ")
    75  	for {
    76  		if retry%5 == 4 {
    77  			fmt.Fprintf(os.Stdout, "%ds --- ", retry+1)
    78  		}
    79  		primaryURL = tools.GetPrimaryURL()
    80  		err = api.Health(tools.BaseAPIParams(primaryURL), true /*primary is ready to rebalance*/)
    81  		if err == nil {
    82  			fmt.Fprintln(os.Stdout, "")
    83  			break
    84  		}
    85  		if retry >= retryCount {
    86  			fmt.Fprintln(os.Stdout, "")
    87  			err = fmt.Errorf("timed out waiting for cluster startup: %v", err)
    88  			return
    89  		}
    90  		retry++
    91  		time.Sleep(sleep)
    92  	}
    93  	tlog.Logln("Cluster is ready")
    94  	return
    95  }
    96  
    97  func initTestEnv() {
    98  	tools.InitLocalCluster()
    99  	proxyURL = tools.RandomProxyURL()
   100  	baseParams = tools.BaseAPIParams(proxyURL)
   101  }
   102  
   103  func TestMain(m *testing.M) {
   104  	flag.Parse()
   105  
   106  	var (
   107  		primaryURL string
   108  		err        error
   109  		exists     bool
   110  	)
   111  
   112  	initTestEnv()
   113  	if cliBck, err = setBucket(); err == nil {
   114  		primaryURL, err = waitForCluster()
   115  	}
   116  
   117  	if err != nil {
   118  		goto fail
   119  	}
   120  
   121  	if !cliBck.IsAIS() {
   122  		exists, err = tools.BucketExists(nil, tools.GetPrimaryURL(), cliBck)
   123  		if err == nil && !exists {
   124  			s := "%q not found \n(hint: "
   125  			s += "check whether %q exists and make sure to build aisnode executable with the corresponding build tag)"
   126  			err = fmt.Errorf(s, cliBck, cliBck)
   127  		}
   128  		if err != nil {
   129  			goto fail
   130  		}
   131  	}
   132  
   133  	initialClusterConfig, err = api.GetClusterConfig(tools.BaseAPIParams(primaryURL))
   134  	if err != nil {
   135  		goto fail
   136  	}
   137  
   138  	m.Run()
   139  	return
   140  
   141  fail:
   142  	tlog.Logln("FAIL: " + err.Error())
   143  	os.Exit(1)
   144  }
   145  
   146  func TestInvalidHTTPMethod(t *testing.T) {
   147  	bp := tools.BaseAPIParams()
   148  	proxyURL := tools.RandomProxyURL(t)
   149  
   150  	req, err := http.NewRequest("TEST", proxyURL, http.NoBody)
   151  	tassert.CheckFatal(t, err)
   152  	tassert.DoAndCheckResp(t, bp.Client, req, http.StatusBadRequest)
   153  }