La Vita è Bear

Generic Go Must function for testing code

Because of go’s error handling design, a lot of functions return 2 values: the result and an error. But there are certain cases handling 2 return values are very inconvenient (the most common such case is when you declare global variables), so for a lot of functions that makes sense to be storing the result in a global variable, you often see them also provide a Must- version that only returns the result without the error.

For example, regexp.Compile also comes with regexp.MustCompile, which simply does (as of go1.20.3):

func MustCompile(str string) *Regexp {
	regexp, err := Compile(str)
	if err != nil {
		panic(`regexp: Compile(` + quote(str) + `): ` + err.Error())
	}
	return regexp
}

Basically they just panics when error happens. In general you should avoid panicking in go code, but global variable declaration is one of the cases panicking is appropriate.

Now, with generics finally being added in go1.18, one can trivially write a generic version of Must like this:

func Must[T any](v T, err error) T {
  if err != nil {
    panic(err)
  }
  return v
}

So instead of using var re = regexp.MustCompile(foo), you can also use var re = Must(regexp.Compile(foo)).

But global variable declarations are only one case to use the Must/Must- functions. Another common case is to write unit tests, especially when you use table tests (I’m just continuing using regexp.Compile as an example here, not that it actually makes sense to write an unit test like this):

func TestFoo(t *testing.T) {
  for _, c := range []struct{
    label string
    re    *regexp.Regexp
    ...
  }{
    {...},
  } {
    t.Run(c.label, func(t *testing.T) {
      // do something with c.re and check the results
    })
  }
}

Now, I must first admit that the panicking implementation of the Must/Must- functions are still perfectly fine to be used in those test cases, as in go when an unit test panics it just fails the test, which is exactly what we want. But, failing the test by panicking will produce more verbose and probably less human readable failing logs, so it’s still better to fail the test differently (more natively). So how do we implement a generic version of MustTesting that does not use panic?

You might think that we can just add testing.TB into the args:

func MustTesting[T any](tb testing.TB, v T, err error) T {
  tb.Helper()
  if err != nil {
    tb.Fatal(err)
  }
  return v
}

But this won’t work as expected, as you cannot write code like MustTesting(t, regexp.Compile(foo)). When you use multi-value return from a function as the args of another function, they must match exactly, you cannot add more args on either side.

When you need to add extra args to a function but do not change the function’s signature, the nature way of doing so in go is by using lambdas, so you can write MustTesting in this way:

func MustTesting[T any](tb testing.TB) func(v T, err error) T {
  return func(v T, err error) T {
    tb.Helper()
    if err != nil {
      tb.Fatal(err)
    }
    return v
  }
}

Now this works, but it comes with its own inconvenience: If you try to use it like MustTesting(t)(regexp.Compile(foo)), you’ll get the compiler error of cannot infer T, because T is not used directly in the first level of function call, so the compiler cannot infer the actual type of T, and you must use it in the more verbose way of MustTesting[*regexp.Regexp](t)(regexp.Compile(foo)).

But fear not, we can improve on that! If we just swap the args of the two level of function calls like this:

func MustTesting[T any](v T, err error) func(tb testing.TB) T {
  return func(tb testing.TB) T {
    tb.Helper()
    if err != nil {
      tb.Fatal(err)
    }
    return v
  }
}

Then you can just use it as MustTesting(regexp.Compile(foo))(t) and it just works.

An extra side/footer note: you don’t really need a Must-/Must/MustTesting function, you can just take advantage of lambdas in global variable declarations, like this:

var re = func() *regexp.Regexp {
  re, err := regexp.Compile(foo)
  if err != nil {
    panic(err)
  }
  return re
}()

So all those Must-/Must/MustTesting functions are just language sugars to make your life easier, and this blog post is just about finding the better language sugar 😄

#English #go #tech