deferred_status.h

#ifndef THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_DEFERRED_STATUS_H_
#define THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_DEFERRED_STATUS_H_

#include <functional>
#include <optional>

#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "grpcpp/support/status.h"

namespace milotic {
// Helper for tests that use callback based gRPC calls.
class DeferredStatus {
 public:
  DeferredStatus() = default;

  std::function<void(grpc::Status)> Setter() {
    return [this](grpc::Status status) ABSL_LOCKS_EXCLUDED(mutex_) {
      absl::MutexLock lock(&mutex_);
      status_ = status;
    };
  }

  grpc::Status AwaitStatus(absl::Duration timeout = absl::Seconds(1))
      ABSL_LOCKS_EXCLUDED(mutex_) {
    auto check = [this]() ABSL_SHARED_LOCKS_REQUIRED(mutex_) {
      return status_.has_value();
    };
    absl::MutexLock lock(&mutex_);
    if (!mutex_.AwaitWithTimeout(absl::Condition(&check), timeout)) {
      return grpc::Status(grpc::StatusCode::DEADLINE_EXCEEDED, "TEST: timeout");
    }
    return status_.value();
  }

 private:
  absl::Mutex mutex_;
  std::optional<grpc::Status> status_ ABSL_GUARDED_BY(mutex_);
};
}  // namespace milotic

#endif  // THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_DEFERRED_STATUS_H_