feat: allow taking benchmarking instance in NewAssert (#1607)

Author: ivokub

test/assert.go

@@ -32,32 +32,53 @@ var (
 // Assert is a helper to test circuits
 type Assert struct {
 	t *testing.T
+	b *testing.B
 	*require.Assertions
 }
 
-// NewAssert returns an Assert helper embedding a testify/require object for convenience
+// NewAssert returns an Assert helper embedding a testify/require object for convenience.
+// It accepts either a *testing.T or *testing.B object.
 //
-// The Assert object caches the compiled circuit:
-//
-// the first call to assert.ProverSucceeded/Failed will compile the circuit for n curves, m backends
-// and subsequent calls will re-use the result of the compilation, if available.
-func NewAssert(t *testing.T) *Assert {
-	return &Assert{t: t, Assertions: require.New(t)}
+// The Assert object caches the compiled circuit. This means that the first call
+// to [Assert.CheckCircuit] will compile the circuit for n curves, m
+// backends and subsequent calls will re-use the result of the compilation, if
+// available. Be careful when benchmarking!
+func NewAssert(tb testing.TB) *Assert {
+	switch t := (tb).(type) {
+	case *testing.T:
+		return &Assert{t: t, Assertions: require.New(t)}
+	case *testing.B:
+		return &Assert{b: t, Assertions: require.New(t)}
+	default:
+		panic("unknown testing type")
+	}
 }
 
 // Run runs the test function fn as a subtest. The subtest is parametrized by
 // the description strings descs.
 func (assert *Assert) Run(fn func(assert *Assert), descs ...string) {
 	desc := strings.Join(descs, "/")
-	assert.t.Run(desc, func(t *testing.T) {
-		assert := &Assert{t, require.New(t)}
-		fn(assert)
-	})
+	if assert.b != nil {
+		assert.b.Run(desc, func(b *testing.B) {
+			assert := &Assert{b: b, Assertions: require.New(b)}
+			fn(assert)
+		})
+	} else {
+		assert.t.Run(desc, func(t *testing.T) {
+			assert := &Assert{t: t, Assertions: require.New(t)}
+			fn(assert)
+		})
+	}
 }
 
 // Log logs using the test instance logger.
 func (assert *Assert) Log(v ...interface{}) {
-	assert.t.Log(v...)
+	if assert.b != nil {
+		assert.b.Log(v...)
+		return
+	} else {
+		assert.t.Log(v...)
+	}
 }
 
 // ProverSucceeded is deprecated: use [Assert.CheckCircuit] instead

test/assert_checkcircuit.go

@@ -101,7 +101,9 @@ func (assert *Assert) CheckCircuit(circuit frontend.Circuit, opts ...TestingOpti
 					}
 
 					// we need to run the setup, prove and verify and check serialization
-					assert.t.Parallel()
+					if assert.t != nil {
+						assert.t.Parallel()
+					}
 
 					var concreteBackend tBackend
 

test/assert_solidity.go

@@ -24,7 +24,11 @@ func (assert *Assert) solidityVerification(b backend.ID, vk solidity.VerifyingKe
 	if !SolcCheck || len(validPublicWitness.Vector().(fr_bn254.Vector)) == 0 {
 		return // nothing to check, will make solc fail.
 	}
-	assert.t.Helper()
+	if assert.b != nil {
+		assert.b.Helper()
+	} else {
+		assert.t.Helper()
+	}
 
 	// make temp dir
 	tmpDir, err := os.MkdirTemp("", "gnark-solidity-check*")
@@ -44,7 +48,7 @@ func (assert *Assert) solidityVerification(b backend.ID, vk solidity.VerifyingKe
 	// generate assets
 	// gnark-solidity-checker generate --dir tmpdir --solidity contract_g16.sol
 	cmd := exec.Command("gnark-solidity-checker", "generate", "--dir", tmpDir, "--solidity", "gnark_verifier.sol")
-	assert.t.Log("running ", cmd.String())
+	assert.Log("running ", cmd.String())
 	out, err := cmd.CombinedOutput()
 	assert.NoError(err, string(out))
 
@@ -90,7 +94,7 @@ func (assert *Assert) solidityVerification(b backend.ID, vk solidity.VerifyingKe
 	// verify proof
 	// gnark-solidity-checker verify --dir tmdir --groth16 --nb-public-inputs 1 --proof 1234 --public-inputs dead
 	cmd = exec.Command("gnark-solidity-checker", checkerOpts...)
-	assert.t.Log("running ", cmd.String())
+	assert.Log("running ", cmd.String())
 	out, err = cmd.CombinedOutput()
 	assert.NoError(err, string(out))
 }