Restore mocked variables in GoLang unit test

One of the guarding principles of writing unit test is that there should be no real external calls for dependant services. Unit test should run by its own and can run without issues on any environment including local, build, test environment. This indicates there should be some mock responses whenever an external call is needed so that different unit test scenarios can be covered.

How can this be done in GoLang? In GoLang, anything can be assigned to a variable even including functions. A variable can be created to reference a function and it can be called just like a normal function when needed. an example of such assignment would be.

var getError = func(str string) error {
	return errors.New(str)
}

A variable getError is created and assigned a function. This means getError is a function and it can be called like:

getError("dummy error")

So if there is a function which calls this function and different errors need to be returned to simulate different scenarios, one can just create a mock function with same name in testcase and return different errors, in such case, when the testcase is ran it will call the new mock function as the variable is referring to the new mock function.

The testcase would look like

func TestGetError(t *testing.T) {
	err := errors.New("error here")
	getError = func(str string) error {
		return err
	}

	assert.Equal(t, err, getError("dummy error"))
}

Whenever getError is called, it will return errors.New("error here"). Okay, this is solving the problem, mock is working and it will pass the test.

However, what will happen if another testcase is to test something else and it just needs the getError to retunr whatever it is in the real code but the function is been overwritten by the mocked function? This means the original function associated with the variable name needs to be restored. Hence there needs a way to remember the original function definition. The solution might looks like:

func TestGetError(t *testing.T) {
	originalGetError := getError
	err := errors.New("error here")
	getError = func(str string) error {
		return err
	}

	assert.Equal(t, err, getError("dummy error"))
	
	getError = originalGetError
}

This works but it si not that beautiful as every mocked variable needs to be restored with above logic. Is there any more elegant way so that there is no need to write such logic for every mocked variable? Yes, the solution is to have a common function which can handle this automatically. To achieve this, reflect can be used.

reflect provides a hacky way to operate variables at a lower language/memory level, it can be used to access internals of variables like fields and functions without knowing definitions of the variable beforeahead. It provides some ways for developers to check the types and values of any element and hence making decisions based on that. 

The challenges above for restoring all mocked variables are developers must know what the mocked variable type and value is so that they know how they can be restored(although they still need to know them so that they know how to mock them). Question here is there is no need to know them at all to restore them. One just needs to have a reference to the original varibale value and get it assigned back when the test is done, the job is done. 

By using reflect, one can abstract the restore logic as:

func RestoreMock(args ...interface{}) func() {
	origArgs := map[interface{}]reflect.Value{}

	for _, name := range args {
		value := reflect.ValueOf(name)
		if value.Kind() != reflect.Ptr {
			panic("unsupported value")
		}

		pv := reflect.New(value.Type().Elem())
		pv.Elem().Set(reflect.Indirect(value))
		origArgs[name] = pv
	}

	return func() {
		for key, value := range origArgs {
			reflect.ValueOf(key).Elem().Set(value.Elem())
		}
	}
}

How it works can be described below:

1. It takes variadic variables to be restored, they must be pointers to the variables
2. The variables are checked one by one and their value is first retrieved
3. The kind/type of each value is retrieved and if it's not a pointer, a panic will happen
4. Thereafter a new element with the same type and same value content is created and saved(this is basically a deep copy about the value content but not the reference)
5. A function will be returned so that it can be called when the test is done to restore the original values. It just loops through the map and reset with the stored values

The usage of the function would be like:

func TestGetError(t *testing.T) {
	defer RestoreMock(&getError)()

	err := errors.New("error here")
	getError = func(str string) error {
		return err
	}

	assert.Equal(t, err, getError("dummy error"))
}

func TestGetErrorAgain(t *testing.T) {
	errDummy := errors.New("dummy error")
	assert.Equal(t, errDummy, getError("dummy error"))
}

At the beginning of the testcase which needs to mock the function, a defered function will be called so that it can remember the old values and then restore the old values when test is done.

A better way to write the RestoreMock function is

// RestoreMock restores the mock functions
func RestoreMock(args map[interface{}]reflect.Value) {
	for key, value := range args {
		reflect.ValueOf(key).Elem().Set(value.Elem())
	}
}

func SaveValues(args ...interface{}) map[interface{}]reflect.Value {
	origArgs := map[interface{}]reflect.Value{}

	for _, name := range args {
		value := reflect.ValueOf(name)
		if value.Kind() != reflect.Ptr {
			panic("unsupported value")
		}

		pv := reflect.New(value.Type().Elem()) // create a reflect.Value of type original one.
		pv.Elem().Set(reflect.Indirect(value))
		origArgs[name] = pv
	}

	return origArgs
}

With this, the test case can be like:

func TestGetError(t *testing.T) {
	defer RestoreMock(SaveValues(&getError))

	err := errors.New("error here")
	getError = func(str string) error {
		return err
	}

	assert.Equal(t, err, getError("dummy error"))
}

func TestGetErrorAgain(t *testing.T) {
	errDummy := errors.New("dummy error")
	assert.Equal(t, errDummy, getError("dummy error"))
}

This avoids the problem of forgetting to write the trailing () as in the previous example.

Hope this helps those who are struggling with mock functions and trying to find how to restore them elegantly.