github.com/ferranbt/nomad@v0.9.3-0.20190607002617-85c449b7667c/testutil/wait.go (about)

     1  package testutil
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"time"
     7  
     8  	"github.com/hashicorp/nomad/nomad/structs"
     9  	testing "github.com/mitchellh/go-testing-interface"
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  type testFn func() (bool, error)
    14  type errorFn func(error)
    15  
    16  func WaitForResult(test testFn, error errorFn) {
    17  	WaitForResultRetries(500*TestMultiplier(), test, error)
    18  }
    19  
    20  func WaitForResultRetries(retries int64, test testFn, error errorFn) {
    21  	for retries > 0 {
    22  		time.Sleep(10 * time.Millisecond)
    23  		retries--
    24  
    25  		success, err := test()
    26  		if success {
    27  			return
    28  		}
    29  
    30  		if retries == 0 {
    31  			error(err)
    32  		}
    33  	}
    34  }
    35  
    36  // AssertUntil asserts the test function passes throughout the given duration.
    37  // Otherwise error is called on failure.
    38  func AssertUntil(until time.Duration, test testFn, error errorFn) {
    39  	deadline := time.Now().Add(until)
    40  	for time.Now().Before(deadline) {
    41  		success, err := test()
    42  		if !success {
    43  			error(err)
    44  			return
    45  		}
    46  		// Sleep some arbitrary fraction of the deadline
    47  		time.Sleep(until / 30)
    48  	}
    49  }
    50  
    51  // TestMultiplier returns a multiplier for retries and waits given environment
    52  // the tests are being run under.
    53  func TestMultiplier() int64 {
    54  	if IsCI() {
    55  		return 4
    56  	}
    57  
    58  	return 1
    59  }
    60  
    61  // Timeout takes the desired timeout and increases it if running in Travis
    62  func Timeout(original time.Duration) time.Duration {
    63  	return original * time.Duration(TestMultiplier())
    64  }
    65  
    66  func IsCI() bool {
    67  	_, ok := os.LookupEnv("CI")
    68  	return ok
    69  }
    70  
    71  func IsTravis() bool {
    72  	_, ok := os.LookupEnv("TRAVIS")
    73  	return ok
    74  }
    75  
    76  func IsAppVeyor() bool {
    77  	_, ok := os.LookupEnv("APPVEYOR")
    78  	return ok
    79  }
    80  
    81  type rpcFn func(string, interface{}, interface{}) error
    82  
    83  // WaitForLeader blocks until a leader is elected.
    84  func WaitForLeader(t testing.T, rpc rpcFn) {
    85  	WaitForResult(func() (bool, error) {
    86  		args := &structs.GenericRequest{}
    87  		var leader string
    88  		err := rpc("Status.Leader", args, &leader)
    89  		return leader != "", err
    90  	}, func(err error) {
    91  		t.Fatalf("failed to find leader: %v", err)
    92  	})
    93  }
    94  
    95  func RegisterJobWithToken(t testing.T, rpc rpcFn, job *structs.Job, token string) {
    96  	WaitForResult(func() (bool, error) {
    97  		args := &structs.JobRegisterRequest{}
    98  		args.Job = job
    99  		args.WriteRequest.Region = "global"
   100  		args.AuthToken = token
   101  		args.Namespace = structs.DefaultNamespace
   102  		var jobResp structs.JobRegisterResponse
   103  		err := rpc("Job.Register", args, &jobResp)
   104  		return err == nil, fmt.Errorf("Job.Register error: %v", err)
   105  	}, func(err error) {
   106  		t.Fatalf("error registering job: %v", err)
   107  	})
   108  
   109  	t.Logf("Job %q registered", job.ID)
   110  }
   111  
   112  func RegisterJob(t testing.T, rpc rpcFn, job *structs.Job) {
   113  	RegisterJobWithToken(t, rpc, job, "")
   114  }
   115  
   116  func WaitForRunningWithToken(t testing.T, rpc rpcFn, job *structs.Job, token string) []*structs.AllocListStub {
   117  	RegisterJobWithToken(t, rpc, job, token)
   118  
   119  	var resp structs.JobAllocationsResponse
   120  
   121  	WaitForResult(func() (bool, error) {
   122  		args := &structs.JobSpecificRequest{}
   123  		args.JobID = job.ID
   124  		args.QueryOptions.Region = "global"
   125  		args.AuthToken = token
   126  		args.Namespace = structs.DefaultNamespace
   127  		err := rpc("Job.Allocations", args, &resp)
   128  		if err != nil {
   129  			return false, fmt.Errorf("Job.Allocations error: %v", err)
   130  		}
   131  
   132  		if len(resp.Allocations) == 0 {
   133  			return false, fmt.Errorf("0 allocations")
   134  		}
   135  
   136  		for _, alloc := range resp.Allocations {
   137  			if alloc.ClientStatus == structs.AllocClientStatusPending {
   138  				return false, fmt.Errorf("alloc not running: id=%v tg=%v status=%v",
   139  					alloc.ID, alloc.TaskGroup, alloc.ClientStatus)
   140  			}
   141  		}
   142  
   143  		return true, nil
   144  	}, func(err error) {
   145  		require.NoError(t, err)
   146  	})
   147  
   148  	return resp.Allocations
   149  }
   150  
   151  // WaitForRunning runs a job and blocks until all allocs are out of pending.
   152  func WaitForRunning(t testing.T, rpc rpcFn, job *structs.Job) []*structs.AllocListStub {
   153  	return WaitForRunningWithToken(t, rpc, job, "")
   154  }