test: fix flaky TestErrCtx and TestContextError (#2670)

Author: miparnisari

internal/middleware/usagemetrics/usagemetrics_test.go

@@ -9,7 +9,9 @@ import (
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/testing/testpb"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
 
 	"github.com/authzed/authzed-go/pkg/responsemeta"
 
@@ -111,7 +113,9 @@ func (s *metricsMiddlewareTestSuite) TestErrCtx() {
 
 	// SimpleCtx times out after two seconds
 	_, err := s.Client.PingError(s.SimpleCtx(), &testpb.PingErrorRequest{}, grpc.Trailer(&trailerMD))
-	s.Require().ErrorContains(err, context.DeadlineExceeded.Error())
+
+	// the error may come from the grpc framework (client-side timeout) or from the API itself (it's a race)
+	s.Require().Equal(codes.DeadlineExceeded, status.Code(err))
 
 	// TODO ideally, this test would assert that no error log has been written
 	// but right now we have no way of capturing the logs

pkg/migrate/migrate_test.go

@@ -2,8 +2,8 @@ package migrate
 
 import (
 	"context"
+	"errors"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -46,7 +46,7 @@ type fakeTx struct{}
 
 func TestContextError(t *testing.T) {
 	req := require.New(t)
-	ctx, cancelFunc := context.WithDeadline(t.Context(), time.Now().Add(1*time.Millisecond))
+	ctx, cancelFunc := context.WithCancel(context.Background())
 	m := NewManager[Driver[fakeConnPool, fakeTx], fakeConnPool, fakeTx]()
 
 	err := m.Register("1", "", func(ctx context.Context, conn fakeConnPool) error {
@@ -56,7 +56,7 @@ func TestContextError(t *testing.T) {
 	req.NoError(err)
 
 	err = m.Register("2", "1", func(ctx context.Context, conn fakeConnPool) error {
-		panic("the second migration should never be executed")
+		return errors.New("the second migration should never be executed")
 	}, noTxMigration)
 	req.NoError(err)