diff --git a/implements.go b/implements.go
index f940575..03692b1 100644
--- a/implements.go
+++ b/implements.go
@@ -23,14 +23,14 @@ import (
 	"context"
 	"fmt"
 
-	empty "google.golang.org/protobuf/types/known/emptypb"
-	pb "google.golang.org/genproto/googleapis/firestore/v1"
+	pb "cloud.google.com/go/firestore/apiv1/firestorepb"
 	codes "google.golang.org/grpc/codes"
 	status "google.golang.org/grpc/status"
+	empty "google.golang.org/protobuf/types/known/emptypb"
 )
 
 // GetDocument overrides the FirestoreServer GetDocument method
-func (s *MockServer) GetDocument(ctx context.Context, req *pb.GetDocumentRequest) (*pb.Document, error) {
+func (s *MockServer) GetDocument(_ context.Context, req *pb.GetDocumentRequest) (*pb.Document, error) {
 	res, err := s.popRPC(req)
 	if err != nil {
 		return nil, err
@@ -69,6 +69,25 @@ func (s *MockServer) BatchGetDocuments(req *pb.BatchGetDocumentsRequest, bs pb.F
 	return nil
 }
 
+func (s *MockServer) ListDocuments(ctx context.Context, req *pb.ListDocumentsRequest) (*pb.ListDocumentsResponse, error) {
+	res, err := s.popRPC(req)
+	if err != nil {
+		return nil, err
+	}
+	responses := res.([]interface{})
+	for _, res := range responses {
+		switch res := res.(type) {
+		case *pb.ListDocumentsResponse:
+			return res, nil
+		case error:
+			return nil, res
+		default:
+			panic(fmt.Sprintf("mockfs.ListDocuments: Bad response type: %+v", res))
+		}
+	}
+	return nil, nil
+}
+
 // RunQuery overrides the FirestoreServer RunQuery method
 func (s *MockServer) RunQuery(req *pb.RunQueryRequest, qs pb.Firestore_RunQueryServer) error {
 	res, err := s.popRPC(req)
@@ -92,6 +111,27 @@ func (s *MockServer) RunQuery(req *pb.RunQueryRequest, qs pb.Firestore_RunQueryS
 	return nil
 }
 
+func (s *MockServer) RunAggregationQuery(req *pb.RunAggregationQueryRequest, qs pb.Firestore_RunAggregationQueryServer) error {
+	res, err := s.popRPC(req)
+	if err != nil {
+		return err
+	}
+	responses := res.([]interface{})
+	for _, res := range responses {
+		switch res := res.(type) {
+		case *pb.RunAggregationQueryResponse:
+			if err := qs.Send(res); err != nil {
+				return err
+			}
+		case error:
+			return res
+		default:
+			panic(fmt.Sprintf("mockfs.RunAggregationQuery: Bad response type: %+v", res))
+		}
+	}
+	return nil
+}
+
 // BeginTransaction overrides the FirestoreServer BeginTransaction method
 func (s *MockServer) BeginTransaction(ctx context.Context, req *pb.BeginTransactionRequest) (*pb.BeginTransactionResponse, error) {
 	res, err := s.popRPC(req)
@@ -133,3 +173,11 @@ func (s *MockServer) Listen(stream pb.Firestore_ListenServer) error {
 	}
 	return nil
 }
+
+func (s *MockServer) BatchWrite(_ context.Context, req *pb.BatchWriteRequest) (*pb.BatchWriteResponse, error) {
+	res, err := s.popRPC(req)
+	if err != nil {
+		return nil, err
+	}
+	return res.(*pb.BatchWriteResponse), nil
+}
diff --git a/server.go b/server.go
index 0f48f8b..f793a55 100644
--- a/server.go
+++ b/server.go
@@ -23,10 +23,10 @@ import (
 	"fmt"
 	"sort"
 
-	"github.com/golang/protobuf/proto"
-	errors "github.com/weathersource/go-errors"
+	pb "cloud.google.com/go/firestore/apiv1/firestorepb"
 	gsrv "github.com/weathersource/go-gsrv"
-	pb "google.golang.org/genproto/googleapis/firestore/v1"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
 )
 
 // MockServer mocks the pb.FirestoreServer interface
@@ -88,9 +88,7 @@ func (s *MockServer) popRPC(gotReq proto.Message) (interface{}, error) {
 		panic("mockfs.popRPC: Out of RPCs.")
 	}
 	ri := s.reqItems[0]
-	resp := s.resps[0]
 	s.reqItems = s.reqItems[1:]
-	s.resps = s.resps[1:]
 	if ri.wantReq != nil {
 		if ri.adjust != nil {
 			ri.adjust(gotReq)
@@ -105,15 +103,18 @@ func (s *MockServer) popRPC(gotReq proto.Message) (interface{}, error) {
 				case *pb.Write_Transform:
 					sort.Sort(byFieldPath(opTyped.Transform.FieldTransforms))
 				}
+				sort.Sort(byFieldPath(w.UpdateTransforms))
 			}
 		}
 
 		if !proto.Equal(gotReq, ri.wantReq) {
-			return nil, errors.NewInternalError(fmt.Sprintf("mockfs.popRPC: Bad request\ngot:  %T\n%s\nwant: %T\n%s",
-				gotReq, proto.MarshalTextString(gotReq),
-				ri.wantReq, proto.MarshalTextString(ri.wantReq)))
+			return nil, fmt.Errorf("mockServer: bad request\ngot:\n%T\n%s\nwant:\n%T\n%s",
+				gotReq, prototext.Format(gotReq),
+				ri.wantReq, prototext.Format(ri.wantReq))
 		}
 	}
+	resp := s.resps[0]
+	s.resps = s.resps[1:]
 	if err, ok := resp.(error); ok {
 		return nil, err
 	}
diff --git a/server_test.go b/server_test.go
index 81d03ac..70c2326 100644
--- a/server_test.go
+++ b/server_test.go
@@ -4,10 +4,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/protobuf/proto"
-	tspb "google.golang.org/protobuf/types/known/timestamppb"
+	pb "cloud.google.com/go/firestore/apiv1/firestorepb"
 	assert "github.com/stretchr/testify/assert"
-	pb "google.golang.org/genproto/googleapis/firestore/v1"
+	"google.golang.org/protobuf/proto"
+	tspb "google.golang.org/protobuf/types/known/timestamppb"
 )
 
 var (