github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/e2e/e2eutil/allocs.go (about)

     1  package e2eutil
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"reflect"
     7  	"strings"
     8  	"time"
     9  
    10  	api "github.com/hashicorp/nomad/api"
    11  	"github.com/hashicorp/nomad/testutil"
    12  	"github.com/kr/pretty"
    13  )
    14  
    15  // AllocsByName sorts allocs by Name
    16  type AllocsByName []*api.AllocationListStub
    17  
    18  func (a AllocsByName) Len() int {
    19  	return len(a)
    20  }
    21  
    22  func (a AllocsByName) Less(i, j int) bool {
    23  	return a[i].Name < a[j].Name
    24  }
    25  
    26  func (a AllocsByName) Swap(i, j int) {
    27  	a[i], a[j] = a[j], a[i]
    28  }
    29  
    30  // WaitForAllocStatusExpected polls 'nomad job status' and exactly compares
    31  // the status of all allocations (including any previous versions) against the
    32  // expected list.
    33  func WaitForAllocStatusExpected(jobID, ns string, expected []string) error {
    34  	err := WaitForAllocStatusComparison(
    35  		func() ([]string, error) { return AllocStatuses(jobID, ns) },
    36  		func(got []string) bool { return reflect.DeepEqual(got, expected) },
    37  		nil,
    38  	)
    39  	if err != nil {
    40  		allocs, _ := AllocsForJob(jobID, ns)
    41  		err = fmt.Errorf("%v\nallocs: %v", err, pretty.Sprint(allocs))
    42  	}
    43  	return err
    44  }
    45  
    46  // WaitForAllocStatusComparison is a convenience wrapper that polls the query
    47  // function until the comparison function returns true.
    48  func WaitForAllocStatusComparison(query func() ([]string, error), comparison func([]string) bool, wc *WaitConfig) error {
    49  	var got []string
    50  	var err error
    51  	interval, retries := wc.OrDefault()
    52  	testutil.WaitForResultRetries(retries, func() (bool, error) {
    53  		time.Sleep(interval)
    54  		got, err = query()
    55  		if err != nil {
    56  			return false, err
    57  		}
    58  		return comparison(got), nil
    59  	}, func(e error) {
    60  		err = fmt.Errorf("alloc status check failed: got %#v", got)
    61  	})
    62  	return err
    63  }
    64  
    65  // AllocsForJob returns a slice of key->value maps, each describing the values
    66  // of the 'nomad job status' Allocations section (not actual
    67  // structs.Allocation objects, query the API if you want those)
    68  func AllocsForJob(jobID, ns string) ([]map[string]string, error) {
    69  
    70  	var nsArg = []string{}
    71  	if ns != "" {
    72  		nsArg = []string{"-namespace", ns}
    73  	}
    74  
    75  	cmd := []string{"nomad", "job", "status"}
    76  	params := []string{"-verbose", "-all-allocs", jobID}
    77  	cmd = append(cmd, nsArg...)
    78  	cmd = append(cmd, params...)
    79  
    80  	out, err := Command(cmd[0], cmd[1:]...)
    81  	if err != nil {
    82  		return nil, fmt.Errorf("'nomad job status' failed: %w", err)
    83  	}
    84  
    85  	section, err := GetSection(out, "Allocations")
    86  	if err != nil {
    87  		return nil, fmt.Errorf("could not find Allocations section: %w", err)
    88  	}
    89  
    90  	allocs, err := ParseColumns(section)
    91  	if err != nil {
    92  		return nil, fmt.Errorf("could not parse Allocations section: %w", err)
    93  	}
    94  	return allocs, nil
    95  }
    96  
    97  // AllocTaskEventsForJob returns a map of allocation IDs containing a map of
    98  // Task Event key value pairs
    99  func AllocTaskEventsForJob(jobID, ns string) (map[string][]map[string]string, error) {
   100  	allocs, err := AllocsForJob(jobID, ns)
   101  	if err != nil {
   102  		return nil, err
   103  	}
   104  
   105  	results := make(map[string][]map[string]string)
   106  	for _, alloc := range allocs {
   107  		results[alloc["ID"]] = make([]map[string]string, 0)
   108  
   109  		cmd := []string{"nomad", "alloc", "status", alloc["ID"]}
   110  		out, err := Command(cmd[0], cmd[1:]...)
   111  		if err != nil {
   112  			return nil, fmt.Errorf("querying alloc status: %w", err)
   113  		}
   114  
   115  		section, err := GetSection(out, "Recent Events:")
   116  		if err != nil {
   117  			return nil, fmt.Errorf("could not find Recent Events section: %w", err)
   118  		}
   119  
   120  		events, err := ParseColumns(section)
   121  		if err != nil {
   122  			return nil, fmt.Errorf("could not parse recent events section: %w", err)
   123  		}
   124  		results[alloc["ID"]] = events
   125  	}
   126  
   127  	return results, nil
   128  }
   129  
   130  // AllocsForNode returns a slice of key->value maps, each describing the values
   131  // of the 'nomad node status' Allocations section (not actual
   132  // structs.Allocation objects, query the API if you want those)
   133  func AllocsForNode(nodeID string) ([]map[string]string, error) {
   134  
   135  	out, err := Command("nomad", "node", "status", "-verbose", nodeID)
   136  	if err != nil {
   137  		return nil, fmt.Errorf("'nomad node status' failed: %w", err)
   138  	}
   139  
   140  	section, err := GetSection(out, "Allocations")
   141  	if err != nil {
   142  		return nil, fmt.Errorf("could not find Allocations section: %w", err)
   143  	}
   144  
   145  	allocs, err := ParseColumns(section)
   146  	if err != nil {
   147  		return nil, fmt.Errorf("could not parse Allocations section: %w", err)
   148  	}
   149  	return allocs, nil
   150  }
   151  
   152  // AllocStatuses returns a slice of client statuses
   153  func AllocStatuses(jobID, ns string) ([]string, error) {
   154  
   155  	allocs, err := AllocsForJob(jobID, ns)
   156  	if err != nil {
   157  		return nil, err
   158  	}
   159  
   160  	statuses := []string{}
   161  	for _, alloc := range allocs {
   162  		statuses = append(statuses, alloc["Status"])
   163  	}
   164  	return statuses, nil
   165  }
   166  
   167  // AllocStatusesRescheduled is a helper function that pulls
   168  // out client statuses only from rescheduled allocs.
   169  func AllocStatusesRescheduled(jobID, ns string) ([]string, error) {
   170  
   171  	var nsArg = []string{}
   172  	if ns != "" {
   173  		nsArg = []string{"-namespace", ns}
   174  	}
   175  
   176  	cmd := []string{"nomad", "job", "status"}
   177  	params := []string{"-verbose", jobID}
   178  	cmd = append(cmd, nsArg...)
   179  	cmd = append(cmd, params...)
   180  
   181  	out, err := Command(cmd[0], cmd[1:]...)
   182  	if err != nil {
   183  		return nil, fmt.Errorf("nomad job status failed: %w", err)
   184  	}
   185  
   186  	section, err := GetSection(out, "Allocations")
   187  	if err != nil {
   188  		return nil, fmt.Errorf("could not find Allocations section: %w", err)
   189  	}
   190  
   191  	allocs, err := ParseColumns(section)
   192  	if err != nil {
   193  		return nil, fmt.Errorf("could not parse Allocations section: %w", err)
   194  	}
   195  
   196  	statuses := []string{}
   197  	for _, alloc := range allocs {
   198  
   199  		allocID := alloc["ID"]
   200  
   201  		cmd := []string{"nomad", "alloc", "status"}
   202  		params := []string{"-json", allocID}
   203  		cmd = append(cmd, nsArg...)
   204  		cmd = append(cmd, params...)
   205  
   206  		// reschedule tracker isn't exposed in the normal CLI output
   207  		out, err := Command(cmd[0], cmd[1:]...)
   208  		if err != nil {
   209  			return nil, fmt.Errorf("nomad alloc status failed: %w", err)
   210  		}
   211  
   212  		dec := json.NewDecoder(strings.NewReader(out))
   213  		alloc := &api.Allocation{}
   214  		err = dec.Decode(alloc)
   215  		if err != nil {
   216  			return nil, fmt.Errorf("could not decode alloc status JSON: %w", err)
   217  		}
   218  
   219  		if (alloc.RescheduleTracker != nil &&
   220  			len(alloc.RescheduleTracker.Events) > 0) || alloc.FollowupEvalID != "" {
   221  			statuses = append(statuses, alloc.ClientStatus)
   222  		}
   223  	}
   224  	return statuses, nil
   225  }
   226  
   227  type LogStream int
   228  
   229  const (
   230  	LogsStdErr LogStream = iota
   231  	LogsStdOut
   232  )
   233  
   234  func AllocLogs(allocID string, logStream LogStream) (string, error) {
   235  	cmd := []string{"nomad", "alloc", "logs"}
   236  	if logStream == LogsStdErr {
   237  		cmd = append(cmd, "-stderr")
   238  	}
   239  	cmd = append(cmd, allocID)
   240  	return Command(cmd[0], cmd[1:]...)
   241  }
   242  
   243  // AllocChecks returns the CLI output from 'nomad alloc checks' on the given
   244  // alloc ID.
   245  func AllocChecks(allocID string) (string, error) {
   246  	cmd := []string{"nomad", "alloc", "checks", allocID}
   247  	return Command(cmd[0], cmd[1:]...)
   248  }
   249  
   250  func AllocTaskLogs(allocID, task string, logStream LogStream) (string, error) {
   251  	cmd := []string{"nomad", "alloc", "logs"}
   252  	if logStream == LogsStdErr {
   253  		cmd = append(cmd, "-stderr")
   254  	}
   255  	cmd = append(cmd, allocID, task)
   256  	return Command(cmd[0], cmd[1:]...)
   257  }
   258  
   259  // AllocExec is a convenience wrapper that runs 'nomad alloc exec' with the
   260  // passed execCmd via '/bin/sh -c', retrying if the task isn't ready
   261  func AllocExec(allocID, taskID, execCmd, ns string, wc *WaitConfig) (string, error) {
   262  	var got string
   263  	var err error
   264  	interval, retries := wc.OrDefault()
   265  
   266  	var nsArg = []string{}
   267  	if ns != "" {
   268  		nsArg = []string{"-namespace", ns}
   269  	}
   270  
   271  	cmd := []string{"nomad", "exec"}
   272  	params := []string{"-task", taskID, allocID, "/bin/sh", "-c", execCmd}
   273  	cmd = append(cmd, nsArg...)
   274  	cmd = append(cmd, params...)
   275  
   276  	testutil.WaitForResultRetries(retries, func() (bool, error) {
   277  		time.Sleep(interval)
   278  		got, err = Command(cmd[0], cmd[1:]...)
   279  		return err == nil, err
   280  	}, func(e error) {
   281  		err = fmt.Errorf("exec failed: '%s': %v\nGot: %v", strings.Join(cmd, " "), e, got)
   282  	})
   283  	return got, err
   284  }
   285  
   286  // WaitForAllocFile is a helper that grabs a file via alloc fs and tests its
   287  // contents; useful for checking the results of rendered templates
   288  func WaitForAllocFile(allocID, path string, test func(string) bool, wc *WaitConfig) error {
   289  	var err error
   290  	var out string
   291  	interval, retries := wc.OrDefault()
   292  
   293  	testutil.WaitForResultRetries(retries, func() (bool, error) {
   294  		time.Sleep(interval)
   295  		out, err = Command("nomad", "alloc", "fs", allocID, path)
   296  		if err != nil {
   297  			return false, fmt.Errorf("could not get file %q from allocation %q: %v",
   298  				path, allocID, err)
   299  		}
   300  		return test(out), nil
   301  	}, func(e error) {
   302  		err = fmt.Errorf("test for file content failed: got %#v\nerror: %v", out, e)
   303  	})
   304  	return err
   305  }