mock_ssh_client.h

#ifndef THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_MOCK_SSH_CLIENT_H_
#define THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_MOCK_SSH_CLIENT_H_

#include <cstddef>
#include <memory>
#include <string>
#include <utility>

#include "gmock.h"
#include "absl/base/no_destructor.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "ssh_client.h"

namespace milotic {
class MockExecutionContext : public ExecutionContext {
 public:
  MockExecutionContext() = default;
  MOCK_METHOD(absl::Status, Write, (std::string data), (override));
  MOCK_METHOD(size_t, WriteBytesPending, (), (const, override));
};

class MockExecutionCallbacks : public ExecutionCallbacks {
 public:
  MockExecutionCallbacks() {
    ON_CALL(*this, OnStdout).WillByDefault(testing::Return(absl::OkStatus()));
    ON_CALL(*this, OnStderr).WillByDefault(testing::Return(absl::OkStatus()));
  }
  MOCK_METHOD(absl::Status, OnStdout,
              (milotic::ExecutionContext * context, absl::string_view data),
              (override));

  MOCK_METHOD(absl::Status, OnStderr,
              (milotic::ExecutionContext * context, absl::string_view data),
              (override));
  MOCK_METHOD(void, OnDone,
              (milotic::ExecutionContext * context, absl::StatusOr<int> result),
              (override));
  MOCK_METHOD(absl::Status, OnWritesDone,
              (milotic::ExecutionContext * context, size_t write_bytes_pending),
              (override));
};

class MockSshClient : public SshClient {
 public:
  MOCK_METHOD(absl::StatusOr<std::unique_ptr<ExecutionContext>>, ExecuteCommand,
              (absl::string_view command, ExecutionCallbacks* callbacks,
               const CommandOptions& options),
              (override));
  MOCK_METHOD(bool, IsConnected, (), (const, override));
  static absl::AnyInvocable<CreateFn>& MockCreate() {
    static absl::NoDestructor<absl::AnyInvocable<CreateFn>> mock_create;
    return *mock_create;
  }
  static absl::StatusOr<std::unique_ptr<SshClient>> Create(
      CreationOptions options) {
    CHECK(MockCreate());
    return MockCreate()(std::move(options));
  }
};

}  // namespace milotic

#endif  // THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_MOCK_SSH_CLIENT_H_