// Copyright 2019 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef REMOTING_BASE_GRPC_SUPPORT_GRPC_TEST_UTIL_H_
#define REMOTING_BASE_GRPC_SUPPORT_GRPC_TEST_UTIL_H_

#include <memory>

#include "base/bind.h"
#include "base/location.h"
#include "base/logging.h"
#include "base/macros.h"
#include "remoting/base/grpc_support/grpc_async_executor.h"
#include "third_party/grpc/src/include/grpcpp/grpcpp.h"

namespace base {
class RunLoop;
}  // namespace base

namespace remoting {
namespace test {

// Block and wait until an event is received from the completion queue, and
// check if the tag matches |expected_tag|.
// Returns whether the event is marked "ok".
bool WaitForCompletion(const base::Location& from_here,
                       grpc::CompletionQueue* completion_queue,
                       void* expected_tag);
void WaitForCompletionAndAssertOk(const base::Location& from_here,
                                  grpc::CompletionQueue* completion_queue,
                                  void* expected_tag);

base::OnceCallback<void(const grpc::Status&)>
CheckStatusThenQuitRunLoopCallback(const base::Location& from_here,
                                   grpc::StatusCode expected_status_code,
                                   base::RunLoop* run_loop);

// Helper class for responding to an async server request.
template <typename ResponseType>
class GrpcServerResponder {
 public:
  explicit GrpcServerResponder(grpc::ServerCompletionQueue* completion_queue) {
    completion_queue_ = completion_queue;
  }

  ~GrpcServerResponder() = default;

  bool Respond(const ResponseType& response, const grpc::Status& status) {
    writer_.Finish(response, status, this);
    return WaitForCompletion(FROM_HERE, completion_queue_, this);
  }

  grpc::ServerContext* context() { return &context_; }

  grpc::ServerAsyncResponseWriter<ResponseType>* writer() { return &writer_; }

 private:
  grpc::ServerContext context_;
  grpc::ServerCompletionQueue* completion_queue_ = nullptr;
  grpc::ServerAsyncResponseWriter<ResponseType> writer_{&context_};

  DISALLOW_COPY_AND_ASSIGN(GrpcServerResponder);
};

// Helper class for responding to an async server stream request.
template <typename ResponseType>
class GrpcServerStreamResponder {
 public:
  explicit GrpcServerStreamResponder(
      grpc::ServerCompletionQueue* completion_queue) {
    completion_queue_ = completion_queue;
  }

  ~GrpcServerStreamResponder() { Close(grpc::Status::OK); }

  // Must call WaitForSendMessageResult() once the client has received the
  // message.
  void SendMessage(const ResponseType& response) {
    writer_.Write(response, /* event_tag */ this);
  }

  // Returns true if the client requests for more messages, false if the client
  // has stopped the stream.
  bool WaitForSendMessageResult() {
    return WaitForCompletion(FROM_HERE, completion_queue_, this);
  }

  void Close(const grpc::Status& status) {
    if (closed_) {
      return;
    }
    writer_.Finish(status, /* event_tag */ this);
    bool ok = WaitForCompletion(FROM_HERE, completion_queue_, this);
    if (!ok) {
      LOG(WARNING) << "Failed to finish stream. Connection might be dropped.";
    }
    closed_ = true;
  }

  grpc::ServerContext* context() { return &context_; }

  grpc::ServerAsyncWriter<ResponseType>* writer() { return &writer_; }

 private:
  grpc::ServerContext context_;
  grpc::ServerCompletionQueue* completion_queue_ = nullptr;
  grpc::ServerAsyncWriter<ResponseType> writer_{&context_};
  bool closed_ = false;

  DISALLOW_COPY_AND_ASSIGN(GrpcServerStreamResponder);
};

}  // namespace test
}  // namespace remoting

#endif  // REMOTING_BASE_GRPC_SUPPORT_GRPC_TEST_UTIL_H_