diff --git a/server/grpc_test.go b/server/grpc_test.go
index 843fe2fe1..d24aa0cc3 100644
--- a/server/grpc_test.go
+++ b/server/grpc_test.go
@@ -977,7 +977,8 @@ func TestQuerySQL(t *testing.T) {
 
 func TestQuerySQLUnaryWithError(t *testing.T) {
 
-	ctx := context.Background()
+	stream := &MockServerTransportStream{}
+	ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream)
 	gh, tearDownFunc := setUpTestQuerySQLUnary(ctx, t)
 	defer tearDownFunc()
 
@@ -1023,7 +1024,8 @@ func TestCRUDIndexes(t *testing.T) {
 	m := test.RunCommand(t)
 	defer m.Close()
 
-	ctx := context.Background()
+	stream := &MockServerTransportStream{}
+	ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream)
 	gh := server.NewGRPCHandler(m.API)
 
 	t.Run("CreateIndex", func(t *testing.T) {
diff --git a/server/handler_test.go b/server/handler_test.go
index 0ffaf3d0e..834399587 100644
--- a/server/handler_test.go
+++ b/server/handler_test.go
@@ -40,6 +40,7 @@ import (
 	pb "github.com/pilosa/pilosa/v2/proto"
 	"github.com/pilosa/pilosa/v2/server"
 	"github.com/pilosa/pilosa/v2/test"
+	"google.golang.org/grpc"
 )
 
 func TestHandler_PostSchemaCluster(t *testing.T) {
@@ -1517,7 +1518,9 @@ func TestQueryHistory(t *testing.T) {
 	test.Do(t, "POST", cmd.URL()+"/index/i0/field/f0", "")
 
 	gh := server.NewGRPCHandler(cmd.API)
-	_, err = gh.QuerySQLUnary(context.Background(), &pb.QuerySQLRequest{
+	stream := &MockServerTransportStream{}
+	ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream)
+	_, err = gh.QuerySQLUnary(ctx, &pb.QuerySQLRequest{
 		Sql: `select * from i0`,
 	})