Golang jobserver.JobServer类代码示例

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
}