func checkChildWorkUnits(t *testing.T, j *jobserver.JobServer, parent, child, workSpecName string, expected map[string]map[string]interface{}) {
missing := make(map[string]struct{})
for name := range expected {
missing[name] = struct{}{}
}
units, msg, err := j.GetChildWorkUnits(parent)
if !assert.NoError(t, err) {
return
}
assert.Empty(t, msg)
assert.Len(t, units, 1)
if assert.Contains(t, units, child) {
for _, unit := range units[child] {
assert.Equal(t, child, unit["worker_id"])
assert.Equal(t, workSpecName, unit["work_spec_name"])
if assert.IsType(t, []byte{}, unit["work_unit_key"]) {
bName := unit["work_unit_key"].([]byte)
name := string(bName)
if assert.Contains(t, expected, name) {
assert.Equal(t, expected[name], unit["work_unit_data"])
}
assert.Contains(t, missing, name, "duplicate child work unit")
delete(missing, name)
}
}
}
assert.Empty(t, missing)
}
func getOneWork(t *testing.T, j *jobserver.JobServer) (ok bool, workSpecName, workUnitKey string, workUnitData map[string]interface{}) {
anything, msg, err := j.GetWork("test", map[string]interface{}{"available_gb": 1})
if !assert.NoError(t, err) {
return
}
assert.Empty(t, msg)
// Since we didn't request multiple work units we should always
// get at most one, but maybe none
if assert.NotNil(t, anything) &&
assert.IsType(t, cborrpc.PythonTuple{}, anything) {
tuple := anything.(cborrpc.PythonTuple)
if assert.Len(t, tuple.Items, 3) {
// "no work unit" gets returned as tuple (nil,
// nil, nil)
if tuple.Items[0] != nil &&
assert.IsType(t, "", tuple.Items[0]) &&
assert.IsType(t, []byte{}, tuple.Items[1]) &&
assert.IsType(t, map[string]interface{}{}, tuple.Items[2]) {
ok = true
workSpecName = tuple.Items[0].(string)
bWorkUnitKey := tuple.Items[1].([]byte)
workUnitKey = string(bWorkUnitKey)
workUnitData = tuple.Items[2].(map[string]interface{})
}
}
}
return
}
// listWorkUnits calls GetWorkUnits (as the similarly-named Python
// function does) and validates that the response matches an expected
// set of work units.
func listWorkUnits(t *testing.T, j *jobserver.JobServer, workSpecName string, options map[string]interface{}, expected map[string]map[string]interface{}) {
gwu, msg, err := j.GetWorkUnits(workSpecName, options)
if !assert.NoError(t, err) {
return
}
assert.Empty(t, msg)
missing := make(map[string]struct{})
for name := range expected {
missing[name] = struct{}{}
}
for _, item := range gwu {
if !assert.IsType(t, cborrpc.PythonTuple{}, item) {
continue
}
tuple := item.(cborrpc.PythonTuple)
if !assert.Len(t, tuple.Items, 2) {
continue
}
if !assert.IsType(t, []byte{}, tuple.Items[0]) {
continue
}
name := string(tuple.Items[0].([]byte))
assert.IsType(t, map[string]interface{}{}, tuple.Items[1])
if assert.Contains(t, expected, name, "unexpected work unit") {
assert.Equal(t, expected[name], tuple.Items[1])
}
assert.Contains(t, missing, name, "duplicate work unit")
delete(missing, name)
}
}
// checkWorkUnitStatus makes a weak assertion about a specific work
// unit's status by calling GetWorkUnitStatus for it.
func checkWorkUnitStatus(t *testing.T, j *jobserver.JobServer, workSpecName, workUnitKey string, status jobserver.WorkUnitStatus) {
dicts, msg, err := j.GetWorkUnitStatus(workSpecName, []string{workUnitKey})
if assert.NoError(t, err) {
assert.Empty(t, msg)
if assert.Len(t, dicts, 1) {
assert.Equal(t, status, dicts[0]["status"])
}
}
}
// addWorkUnit packages a single work unit key and data dictionary
// into the tuple format JobServer expects, and calls AddWorkUnits(),
// checking the result.
func addWorkUnit(t *testing.T, j *jobserver.JobServer, workSpecName, key string, data map[string]interface{}) {
keyDataPair := cborrpc.PythonTuple{Items: []interface{}{key, data}}
keyDataList := []interface{}{keyDataPair}
ok, msg, err := j.AddWorkUnits(workSpecName, keyDataList)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
}
// setWorkSpec calls the eponymous JobServer function, checking that
// it ran successfully, and returns the work spec name on success.
func setWorkSpec(t *testing.T, j *jobserver.JobServer, workSpec map[string]interface{}) string {
ok, msg, err := j.SetWorkSpec(workSpec)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
workSpecName, ok := workSpec["name"].(string)
assert.True(t, ok, "workSpec[\"name\"] not a string")
return workSpecName
}
// finishWorkUnit marks a specific work unit as finished.
func finishWorkUnit(t *testing.T, j *jobserver.JobServer, workSpecName, workUnitKey string, data map[string]interface{}) {
options := map[string]interface{}{
"status": jobserver.Finished,
}
if data != nil {
options["data"] = data
}
ok, msg, err := j.UpdateWorkUnit(workSpecName, workUnitKey, options)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
}
// addPrefixedWorkUnits adds a series of similarly-named work units
// to a work spec. If prefix is "u", it adds count work units named
// u001, u002, .... The work spec dictionaries have a single key "k"
// with values v1, v2, ....
func addPrefixedWorkUnits(t *testing.T, j *jobserver.JobServer, workSpecName, prefix string, count int) {
workUnitKvps := make([]interface{}, count)
for i := range workUnitKvps {
key := fmt.Sprintf("%s%03d", prefix, i+1)
data := map[string]interface{}{"k": fmt.Sprintf("v%v", i+1)}
items := []interface{}{key, data}
workUnitKvps[i] = cborrpc.PythonTuple{Items: items}
}
ok, msg, err := j.AddWorkUnits(workSpecName, workUnitKvps)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
}
// addWorkUnits adds a batch of work units to the system in one call.
func addWorkUnits(t *testing.T, j *jobserver.JobServer, workSpecName string, workUnits map[string]map[string]interface{}) {
// Assemble the parameters to AddWorkUnits as one big list of
// pairs of (key, data)
var awu []interface{}
for name, data := range workUnits {
pair := []interface{}{name, data}
awu = append(awu, pair)
}
ok, msg, err := j.AddWorkUnits(workSpecName, awu)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
}
// delWorkUnitsBy is the core of the DelWorkUnits tests that expect to
// delet single work units. It calls options(state) to get options
// to DelWorkUnits, and verifies that this deletes the single work unit
// associated with state.
func delWorkUnitsBy(t *testing.T, j *jobserver.JobServer, n int, state jobserver.WorkUnitStatus, options func(jobserver.WorkUnitStatus) map[string]interface{}) {
workSpecName, expected := prepareSomeOfEach(t, j, n)
delete(expected, stateShortName[state])
count, msg, err := j.DelWorkUnits(workSpecName, options(state))
if assert.NoError(t, err) {
assert.Equal(t, 1, count)
assert.Empty(t, msg)
}
listWorkUnits(t, j, workSpecName, gwuEverything, expected)
_, err = j.Clear()
assert.NoError(t, err)
}
// prioritizeWorkUnit changes the priority of a single work unit.
func prioritizeWorkUnit(t *testing.T, j *jobserver.JobServer, workSpecName, key string, priority int, adjust bool) {
options := map[string]interface{}{
"work_unit_keys": []interface{}{key},
}
if adjust {
options["priority"] = nil
options["adjustment"] = priority
} else {
options["priority"] = priority
options["adjustment"] = nil
}
ok, msg, err := j.PrioritizeWorkUnits(workSpecName, options)
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
}
// getSpecificWork calls GetWork expecting a specific work unit to
// come back, and returns its data dictionary.
func getSpecificWork(t *testing.T, j *jobserver.JobServer, workSpecName, workUnitKey string) map[string]interface{} {
anything, msg, err := j.GetWork("test", map[string]interface{}{"available_gb": 1})
if !assert.NoError(t, err) {
return nil
}
assert.Empty(t, msg)
if assert.NotNil(t, anything) && assert.IsType(t, cborrpc.PythonTuple{}, anything) {
tuple := anything.(cborrpc.PythonTuple)
if assert.Len(t, tuple.Items, 3) && assert.NotNil(t, tuple.Items[0]) {
assert.Equal(t, workSpecName, tuple.Items[0])
assert.Equal(t, []byte(workUnitKey), tuple.Items[1])
if assert.IsType(t, tuple.Items[2], map[string]interface{}{}) {
return tuple.Items[2].(map[string]interface{})
}
}
}
return nil
}
// getOneWorkUnit calls GetWorkUnits for a single specific work unit,
// checks the results, and returns its data dictionary (or nil if absent).
func getOneWorkUnit(t *testing.T, j *jobserver.JobServer, workSpecName, workUnitKey string) map[string]interface{} {
list, msg, err := j.GetWorkUnits(workSpecName, map[string]interface{}{"work_unit_keys": []interface{}{workUnitKey}})
if !assert.NoError(t, err) {
return nil
}
assert.Empty(t, msg)
if !assert.Len(t, list, 1) {
return nil
}
if !assert.IsType(t, cborrpc.PythonTuple{}, list[0]) {
return nil
}
tuple := list[0].(cborrpc.PythonTuple)
if !assert.Len(t, tuple.Items, 2) {
return nil
}
assert.Equal(t, []byte(workUnitKey), tuple.Items[0])
if assert.IsType(t, map[string]interface{}{}, tuple.Items[1]) {
return tuple.Items[1].(map[string]interface{})
}
return nil
}
func prepareSomeOfEach(t *testing.T, j *jobserver.JobServer, n int) (workSpecName string, expected map[string]map[string]interface{}) {
data := map[string]interface{}{"x": 1}
expected = map[string]map[string]interface{}{}
workSpecName = setWorkSpec(t, j, WorkSpecData)
for _, name := range []string{"FA", "IL"}[:n] {
addWorkUnit(t, j, workSpecName, name, data)
getSpecificWork(t, j, workSpecName, name)
ok, msg, err := j.UpdateWorkUnit(workSpecName, name, map[string]interface{}{"status": jobserver.Failed})
if assert.NoError(t, err) {
assert.True(t, ok)
assert.Empty(t, msg)
}
expected[name] = data
}
for _, name := range []string{"FI", "NI"}[:n] {
addWorkUnit(t, j, workSpecName, name, data)
getSpecificWork(t, j, workSpecName, name)
finishWorkUnit(t, j, workSpecName, name, nil)
expected[name] = data
}
for _, name := range []string{"PE", "ND"}[:n] {
addWorkUnit(t, j, workSpecName, name, data)
getSpecificWork(t, j, workSpecName, name)
expected[name] = data
}
for _, name := range []string{"AV", "AI"}[:n] {
addWorkUnit(t, j, workSpecName, name, data)
expected[name] = data
}
return
}
