gRPC: add more unit tests to Datastore and Stream (#1943)

Author: var-const

Firestore/core/src/firebase/firestore/remote/datastore.h

@@ -72,6 +72,11 @@ class Datastore : public std::enable_shared_from_this<Datastore> {
             auth::CredentialsProvider* credentials,
             FSTSerializerBeta* serializer);
 
+  virtual ~Datastore() {
+  }
+
+  /** Starts polling the gRPC completion queue. */
+  void Start();
   /** Cancels any pending gRPC calls and drains the gRPC completion queue. */
   void Shutdown();
 
@@ -101,6 +106,16 @@ class Datastore : public std::enable_shared_from_this<Datastore> {
   Datastore& operator=(const Datastore& other) = delete;
   Datastore& operator=(Datastore&& other) = delete;
 
+ protected:
+  /** Test-only method */
+  grpc::CompletionQueue* grpc_queue() {
+    return &grpc_queue_;
+  }
+  /** Test-only method */
+  GrpcCall* LastCall() {
+    return !active_calls_.empty() ? active_calls_.back().get() : nullptr;
+  }
+
  private:
   void PollGrpcQueue();
 

Firestore/core/src/firebase/firestore/remote/datastore.mm

@@ -95,6 +95,9 @@ void LogGrpcCallFinished(absl::string_view rpc_name,
       grpc_connection_{database_info, worker_queue, &grpc_queue_,
                        connectivity_monitor_.get()},
       serializer_bridge_{NOT_NULL(serializer)} {
+}
+
+void Datastore::Start() {
   rpc_executor_->Execute([this] { PollGrpcQueue(); });
 }
 

Firestore/core/src/firebase/firestore/remote/grpc_call.h

@@ -20,6 +20,7 @@
 #include <map>
 
 #include "Firestore/core/src/firebase/firestore/util/status.h"
+#include "grpcpp/client_context.h"
 #include "grpcpp/support/string_ref.h"
 
 namespace firebase {
@@ -55,6 +56,9 @@ class GrpcCall {
 
   /** Finishes the call with an error, notifying any callbacks and observers. */
   virtual void FinishAndNotify(const util::Status& status) = 0;
+
+  /** For tests only */
+  virtual grpc::ClientContext* context() = 0;
 };
 
 }  // namespace remote

Firestore/core/src/firebase/firestore/remote/grpc_stream.cc

@@ -98,6 +98,7 @@ GrpcStream::GrpcStream(
 GrpcStream::~GrpcStream() {
   HARD_ASSERT(completions_.empty(),
               "GrpcStream is being destroyed without proper shutdown");
+  MaybeUnregister();
 }
 
 void GrpcStream::Start() {
@@ -172,11 +173,7 @@ void GrpcStream::FinishAndNotify(const Status& status) {
 }
 
 void GrpcStream::Shutdown() {
-  if (grpc_connection_) {
-    grpc_connection_->Unregister(this);
-    grpc_connection_ = nullptr;
-  }
-
+  MaybeUnregister();
   if (completions_.empty()) {
     // Nothing to cancel.
     return;
@@ -202,6 +199,13 @@ void GrpcStream::Shutdown() {
   FastFinishCompletionsBlocking();
 }
 
+void GrpcStream::MaybeUnregister() {
+  if (grpc_connection_) {
+    grpc_connection_->Unregister(this);
+    grpc_connection_ = nullptr;
+  }
+}
+
 void GrpcStream::FastFinishCompletionsBlocking() {
   // TODO(varconst): reset buffered_writer_? Should not be necessary, because it
   // should never be called again after a call to Finish.

Firestore/core/src/firebase/firestore/remote/grpc_stream.h

@@ -173,7 +173,7 @@ class GrpcStream : public GrpcCall {
   Metadata GetResponseHeaders() const override;
 
   /** For tests only */
-  grpc::ClientContext* context() {
+  grpc::ClientContext* context() override {
     return context_.get();
   }
 
@@ -185,6 +185,7 @@ class GrpcStream : public GrpcCall {
   void UnsetObserver() {
     observer_ = nullptr;
   }
+  void MaybeUnregister();
 
   void OnStart();
   void OnRead(const grpc::ByteBuffer& message);

Firestore/core/src/firebase/firestore/remote/grpc_streaming_reader.h

@@ -88,7 +88,7 @@ class GrpcStreamingReader : public GrpcCall, public GrpcStreamObserver {
   void OnStreamFinish(const util::Status& status) override;
 
   /** For tests only */
-  grpc::ClientContext* context() {
+  grpc::ClientContext* context() override {
     return stream_->context();
   }
 

Firestore/core/src/firebase/firestore/remote/grpc_unary_call.cc

@@ -46,6 +46,7 @@ GrpcUnaryCall::GrpcUnaryCall(
 GrpcUnaryCall::~GrpcUnaryCall() {
   HARD_ASSERT(!finish_completion_,
               "GrpcUnaryCall is being destroyed without proper shutdown");
+  MaybeUnregister();
 }
 
 void GrpcUnaryCall::Start(Callback&& callback) {
@@ -58,6 +59,8 @@ void GrpcUnaryCall::Start(Callback&& callback) {
       [this](bool /*ignored_ok*/, const GrpcCompletion* completion) {
         // Ignoring ok, status should contain all the relevant information.
         finish_completion_ = nullptr;
+        Shutdown();
+
         auto callback = std::move(callback_);
         if (completion->status()->ok()) {
           callback(*completion->message());
@@ -84,11 +87,7 @@ void GrpcUnaryCall::FinishAndNotify(const util::Status& status) {
 }
 
 void GrpcUnaryCall::Shutdown() {
-  if (grpc_connection_) {
-    grpc_connection_->Unregister(this);
-    grpc_connection_ = nullptr;
-  }
-
+  MaybeUnregister();
   if (!finish_completion_) {
     // Nothing to cancel.
     return;
@@ -102,6 +101,13 @@ void GrpcUnaryCall::Shutdown() {
   finish_completion_ = nullptr;
 }
 
+void GrpcUnaryCall::MaybeUnregister() {
+  if (grpc_connection_) {
+    grpc_connection_->Unregister(this);
+    grpc_connection_ = nullptr;
+  }
+}
+
 GrpcCall::Metadata GrpcUnaryCall::GetResponseHeaders() const {
   return context_->GetServerInitialMetadata();
 }

Firestore/core/src/firebase/firestore/remote/grpc_unary_call.h

@@ -82,12 +82,13 @@ class GrpcUnaryCall : public GrpcCall {
   Metadata GetResponseHeaders() const override;
 
   /** For tests only */
-  grpc::ClientContext* context() {
+  grpc::ClientContext* context() override {
     return context_.get();
   }
 
  private:
   void Shutdown();
+  void MaybeUnregister();
 
   // See comments in `GrpcStream` on lifetime issues for gRPC objects.
   std::unique_ptr<grpc::ClientContext> context_;

Firestore/core/src/firebase/firestore/remote/remote_objc_bridge.h

@@ -161,7 +161,8 @@ class DatastoreSerializer {
   NSArray<FSTMaybeDocument*>* MergeLookupResponses(
       const std::vector<grpc::ByteBuffer>& responses,
       util::Status* out_status) const;
-  FSTMaybeDocument* ToMaybeDocument(GCFSBatchGetDocumentsResponse*) const;
+  FSTMaybeDocument* ToMaybeDocument(
+      GCFSBatchGetDocumentsResponse* response) const;
 
   FSTSerializerBeta* GetSerializer() {
     return serializer_;

Firestore/core/src/firebase/firestore/remote/remote_objc_bridge.mm

@@ -60,10 +60,15 @@
   return output.str();
 }
 
-NSData* ConvertToNsData(const grpc::ByteBuffer& buffer) {
+NSData* ConvertToNsData(const grpc::ByteBuffer& buffer, NSError** out_error) {
   std::vector<grpc::Slice> slices;
   grpc::Status status = buffer.Dump(&slices);
-  HARD_ASSERT(status.ok(), "Trying to convert an invalid grpc::ByteBuffer");
+  if (!status.ok()) {
+    *out_error =
+        MakeNSError(Status{FirestoreErrorCode::Internal,
+                           "Trying to convert an invalid grpc::ByteBuffer"});
+    return nil;
+  }
 
   if (slices.size() == 1) {
     return [NSData dataWithBytes:slices.front().begin()
@@ -86,10 +91,13 @@
 template <typename Proto>
 Proto* ToProto(const grpc::ByteBuffer& message, Status* out_status) {
   NSError* error = nil;
-  Proto* proto = [Proto parseFromData:ConvertToNsData(message) error:&error];
+  NSData* data = ConvertToNsData(message, &error);
   if (!error) {
-    *out_status = Status::OK();
-    return proto;
+    Proto* proto = [Proto parseFromData:data error:&error];
+    if (!error) {
+      *out_status = Status::OK();
+      return proto;
+    }
   }
 
   std::string error_description = StringFormat(