chore: improve metrics middleware test (#2812)

Author: tstirrat15

internal/middleware/usagemetrics/usagemetrics.go

@@ -12,6 +12,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/authzed/authzed-go/pkg/responsemeta"
+	"github.com/authzed/ctxkey"
 	"github.com/authzed/grpcutil"
 
 	log "github.com/authzed/spicedb/internal/logging"
@@ -87,40 +88,25 @@ func StreamServerInterceptor() grpc.StreamServerInterceptor {
 }
 
 // Create a new type to prevent context collisions
-type responseMetaKey string
-
-var metadataCtxKey responseMetaKey = "dispatched-response-meta"
-
-type metaHandle struct{ metadata *dispatch.ResponseMeta }
+var ctxResponseMeta = ctxkey.NewBoxedWithDefault[*dispatch.ResponseMeta](nil)
 
 // SetInContext should be called in a gRPC handler to correctly set the response metadata
 // for the dispatched request.
 func SetInContext(ctx context.Context, metadata *dispatch.ResponseMeta) {
-	possibleHandle := ctx.Value(metadataCtxKey)
-	if possibleHandle == nil {
-		return
-	}
-
-	handle := possibleHandle.(*metaHandle)
-	handle.metadata = metadata
+	ctxResponseMeta.Set(ctx, metadata)
 }
 
 // FromContext returns any metadata that was stored in the context.
 //
 // This is useful for testing that a handler is properly setting the context.
 func FromContext(ctx context.Context) *dispatch.ResponseMeta {
-	possibleHandle := ctx.Value(metadataCtxKey)
-	if possibleHandle == nil {
-		return nil
-	}
-	return possibleHandle.(*metaHandle).metadata
+	return ctxResponseMeta.Value(ctx)
 }
 
 // contextWithHandle creates a new context with a location to store metadata
 // returned from a dispatched request.
 //
 // This should only be called in middleware or testing functions.
 func contextWithHandle(ctx context.Context) context.Context {
-	var handle metaHandle
-	return context.WithValue(ctx, metadataCtxKey, &handle)
+	return ctxResponseMeta.WithBox()(ctx)
 }

internal/middleware/usagemetrics/usagemetrics_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/testing/testpb"
 	"github.com/stretchr/testify/suite"
@@ -111,8 +112,8 @@ func (s *metricsMiddlewareTestSuite) TestTrailers_Stream() {
 func (s *metricsMiddlewareTestSuite) TestErrCtx() {
 	var trailerMD metadata.MD
 
-	// SimpleCtx times out after two seconds
-	_, err := s.Client.PingError(s.SimpleCtx(), &testpb.PingErrorRequest{}, grpc.Trailer(&trailerMD))
+	// Make a request with an immediate deadline
+	_, err := s.Client.PingError(s.DeadlineCtx(time.Now()), &testpb.PingErrorRequest{}, grpc.Trailer(&trailerMD))
 
 	// 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))