github.com/anuvu/nomad@v0.8.7-atom1/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  	"github.com/mitchellh/go-testing-interface"
    10  )
    11  
    12  const (
    13  	// TravisRunEnv is an environment variable that is set if being run by
    14  	// Travis.
    15  	TravisRunEnv = "CI"
    16  )
    17  
    18  type testFn func() (bool, error)
    19  type errorFn func(error)
    20  
    21  func WaitForResult(test testFn, error errorFn) {
    22  	WaitForResultRetries(500*TestMultiplier(), test, error)
    23  }
    24  
    25  func WaitForResultRetries(retries int64, test testFn, error errorFn) {
    26  	for retries > 0 {
    27  		time.Sleep(10 * time.Millisecond)
    28  		retries--
    29  
    30  		success, err := test()
    31  		if success {
    32  			return
    33  		}
    34  
    35  		if retries == 0 {
    36  			error(err)
    37  		}
    38  	}
    39  }
    40  
    41  // AssertUntil asserts the test function passes throughout the given duration.
    42  // Otherwise error is called on failure.
    43  func AssertUntil(until time.Duration, test testFn, error errorFn) {
    44  	deadline := time.Now().Add(until)
    45  	for time.Now().Before(deadline) {
    46  		success, err := test()
    47  		if !success {
    48  			error(err)
    49  			return
    50  		}
    51  		// Sleep some arbitrary fraction of the deadline
    52  		time.Sleep(until / 30)
    53  	}
    54  }
    55  
    56  // TestMultiplier returns a multiplier for retries and waits given environment
    57  // the tests are being run under.
    58  func TestMultiplier() int64 {
    59  	if IsTravis() {
    60  		return 4
    61  	}
    62  
    63  	return 1
    64  }
    65  
    66  // Timeout takes the desired timeout and increases it if running in Travis
    67  func Timeout(original time.Duration) time.Duration {
    68  	return original * time.Duration(TestMultiplier())
    69  }
    70  
    71  func IsTravis() bool {
    72  	_, ok := os.LookupEnv(TravisRunEnv)
    73  	return ok
    74  }
    75  
    76  type rpcFn func(string, interface{}, interface{}) error
    77  
    78  // WaitForLeader blocks until a leader is elected.
    79  func WaitForLeader(t testing.T, rpc rpcFn) {
    80  	WaitForResult(func() (bool, error) {
    81  		args := &structs.GenericRequest{}
    82  		var leader string
    83  		err := rpc("Status.Leader", args, &leader)
    84  		return leader != "", err
    85  	}, func(err error) {
    86  		t.Fatalf("failed to find leader: %v", err)
    87  	})
    88  }
    89  
    90  // WaitForRunning runs a job and blocks until it is running.
    91  func WaitForRunning(t testing.T, rpc rpcFn, job *structs.Job) {
    92  	registered := false
    93  	WaitForResult(func() (bool, error) {
    94  		if !registered {
    95  			args := &structs.JobRegisterRequest{}
    96  			args.Job = job
    97  			args.WriteRequest.Region = "global"
    98  			var jobResp structs.JobRegisterResponse
    99  			err := rpc("Job.Register", args, &jobResp)
   100  			if err != nil {
   101  				return false, fmt.Errorf("Job.Register error: %v", err)
   102  			}
   103  
   104  			// Only register once
   105  			registered = true
   106  		}
   107  
   108  		args := &structs.JobSummaryRequest{}
   109  		args.JobID = job.ID
   110  		args.QueryOptions.Region = "global"
   111  		var resp structs.JobSummaryResponse
   112  		err := rpc("Job.Summary", args, &resp)
   113  		if err != nil {
   114  			return false, fmt.Errorf("Job.Summary error: %v", err)
   115  		}
   116  
   117  		tgs := len(job.TaskGroups)
   118  		summaries := len(resp.JobSummary.Summary)
   119  		if tgs != summaries {
   120  			return false, fmt.Errorf("task_groups=%d summaries=%d", tgs, summaries)
   121  		}
   122  
   123  		for tg, summary := range resp.JobSummary.Summary {
   124  			if summary.Running == 0 {
   125  				return false, fmt.Errorf("task_group=%s %#v", tg, resp.JobSummary.Summary)
   126  			}
   127  		}
   128  
   129  		return true, nil
   130  	}, func(err error) {
   131  		t.Fatalf("job not running: %v", err)
   132  	})
   133  }