protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

protected:

std::string transport() const override { return "grpc"; }

void SetUp() override { SetUpTest(); }

void TearDown() override { TearDownTest(); }

TestServer server("flight-test-server");

server.Start();

ASSERT_TRUE(server.IsRunning());

std::stringstream ss;

ss << "grpc://localhost:" << server.port();

std::string uri = ss.str();

std::unique_ptr<FlightClient> client;

ASSERT_OK_AND_ASSIGN(auto location1, Location::Parse(uri));

ASSERT_OK_AND_ASSIGN(auto location2, Location::Parse(uri));

ASSERT_OK_AND_ASSIGN(client, FlightClient::Connect(location1));

ASSERT_OK(client->Close());

ASSERT_OK_AND_ASSIGN(client, FlightClient::Connect(location2));

ASSERT_OK(client->Close());

ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("invalid://localhost:1234"));

EXPECT_RAISES_WITH_MESSAGE_THAT(

KeyError, ::testing::HasSubstr("No client transport implementation for invalid"),

FlightClient::Connect(location));

}

#ifndef _WIN32

TEST(TestFlight, ConnectUriUnix) {

TestServer server("flight-test-server", "/tmp/flight-test.sock");

server.Start();

ASSERT_TRUE(server.IsRunning());

std::stringstream ss;

ss << "grpc+unix://" << server.unix_sock();

std::string uri = ss.str();

std::unique_ptr<FlightClient> client;

ASSERT_OK_AND_ASSIGN(auto location1, Location::Parse(uri));

ASSERT_OK_AND_ASSIGN(auto location2, Location::Parse(uri));

ASSERT_OK_AND_ASSIGN(client, FlightClient::Connect(location1));

ASSERT_OK(client->Close());

ASSERT_OK_AND_ASSIGN(client, FlightClient::Connect(location2));

ASSERT_OK(client->Close());

std::unique_ptr<FlightServerBase> server = ExampleTestServer();

ASSERT_OK_AND_ASSIGN(auto location, Location::ForGrpcTcp("[::1]", 0));

FlightServerOptions options(location);

ASSERT_OK(server->Init(options));

ASSERT_GT(server->port(), 0);

ASSERT_OK_AND_ASSIGN(auto location2, Location::ForGrpcTcp("[::1]", server->port()));

std::unique_ptr<FlightClient> client;

ASSERT_OK_AND_ASSIGN(client, FlightClient::Connect(location2));

ASSERT_OK(client->ListFlights());

public:

void SetUp() {

server_ = ExampleTestServer();

ASSERT_OK_AND_ASSIGN(auto location, Location::ForGrpcTcp("localhost", 0));

FlightServerOptions options(location);

ASSERT_OK(server_->Init(options));

ASSERT_OK(ConnectClient());

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

Status ConnectClient() {

ARROW_ASSIGN_OR_RAISE(auto location,

Location::ForGrpcTcp("localhost", server_->port()));

return FlightClient::Connect(location).Value(&client_);

}

template <typename EndpointCheckFunc>

void CheckDoGet(const FlightDescriptor& descr,

const RecordBatchVector& expected_batches,

EndpointCheckFunc&& check_endpoints) {

auto expected_schema = expected_batches[0]->schema();

ASSERT_OK_AND_ASSIGN(auto info, client_->GetFlightInfo(descr));

check_endpoints(info->endpoints());

ipc::DictionaryMemo dict_memo;

ASSERT_OK_AND_ASSIGN(auto schema, info->GetSchema(&dict_memo));

AssertSchemaEqual(*expected_schema, *schema);

// By convention, fetch the first endpoint

Ticket ticket = info->endpoints()[0].ticket;

CheckDoGet(ticket, expected_batches);

}

void CheckDoGet(const Ticket& ticket, const RecordBatchVector& expected_batches) {

auto num_batches = static_cast<int>(expected_batches.size());

ASSERT_GE(num_batches, 2);

ASSERT_OK_AND_ASSIGN(auto stream, client_->DoGet(ticket));

ASSERT_OK_AND_ASSIGN(auto stream2, client_->DoGet(ticket));

ASSERT_OK_AND_ASSIGN(auto reader, MakeRecordBatchReader(std::move(stream2)));

FlightStreamChunk chunk;

std::shared_ptr<RecordBatch> batch;

for (int i = 0; i < num_batches; ++i) {

ASSERT_OK_AND_ASSIGN(chunk, stream->Next());

ASSERT_OK(reader->ReadNext(&batch));

ASSERT_NE(nullptr, chunk.data);

ASSERT_NE(nullptr, batch);

#if !defined(__MINGW32__)

ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk.data);

ASSERT_BATCHES_EQUAL(*expected_batches[i], *batch);

#else

// In MINGW32, the following code does not have the reproducibility at the LSB

// even when this is called twice with the same seed.

// As a workaround, use approxEqual

// /* from GenerateTypedData in random.cc */

// std::default_random_engine rng(seed); // seed = 282475250

// std::uniform_real_distribution<double> dist;

// std::generate(data, data + n, // n = 10

// [&dist, &rng] { return static_cast<ValueType>(dist(rng)); });

// /* data[1] = 0x40852cdfe23d3976 or 0x40852cdfe23d3975 */

ASSERT_BATCHES_APPROX_EQUAL(*expected_batches[i], *chunk.data);

ASSERT_BATCHES_APPROX_EQUAL(*expected_batches[i], *batch);

#endif

}

// Stream exhausted

ASSERT_OK_AND_ASSIGN(chunk, stream->Next());

ASSERT_OK(reader->ReadNext(&batch));

ASSERT_EQ(nullptr, chunk.data);

ASSERT_EQ(nullptr, batch);

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

Status DoAction(const ServerCallContext& context, const Action& action,

std::unique_ptr<ResultStream>* result) override {

auto buf = Buffer::FromString(context.peer_identity());

auto peer = Buffer::FromString(context.peer());

*result = std::make_unique<SimpleResultStream>(

std::vector<Result>{Result{buf}, Result{peer}});

return Status::OK();

}

Status DoAction(const ServerCallContext& context, const Action& action,

std::unique_ptr<ResultStream>* result) override {

auto buf = Buffer::FromString("Hello, world!");

*result = std::make_unique<SimpleResultStream>(std::vector<Result>{Result{buf}});

return Status::OK();

}

public:

Status ListFlights(const ServerCallContext& context, const Criteria* criteria,

std::unique_ptr<FlightListing>* listings) override {

return Status::OK();

}

public:

void SetUp() {

ASSERT_OK(MakeServer<AuthTestServer>(

&server_, &client_,

[](FlightServerOptions* options) {

options->auth_handler =

std::make_unique<TestServerAuthHandler>("user", "p4ssw0rd");

return Status::OK();

},

[](FlightClientOptions* options) { return Status::OK(); }));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

public:

void SetUp() {

ASSERT_OK(MakeServer<AuthTestServer>(

&server_, &client_,

[](FlightServerOptions* options) {

options->auth_handler =

std::make_unique<TestServerBasicAuthHandler>("user", "p4ssw0rd");

return Status::OK();

},

[](FlightClientOptions* options) { return Status::OK(); }));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

public:

void SetUp() {

// Manually initialize gRPC to try to ensure some thread-locals

// get initialized.

// https://github.com/grpc/grpc/issues/13856

// https://github.com/grpc/grpc/issues/20311

// In general, gRPC on MacOS struggles with TLS (both in the sense

// of thread-locals and encryption)

grpc_init();

server_.reset(new TlsTestServer);

ASSERT_OK_AND_ASSIGN(auto location, Location::ForGrpcTls("localhost", 0));

FlightServerOptions options(location);

ASSERT_RAISES(UnknownError, server_->Init(options));

ASSERT_OK(ExampleTlsCertificates(&options.tls_certificates));

ASSERT_OK(server_->Init(options));

ASSERT_OK_AND_ASSIGN(location_, Location::ForGrpcTls("localhost", server_->port()));

ASSERT_OK(ConnectClient());

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

grpc_shutdown();

}

Status ConnectClient() {

auto options = FlightClientOptions::Defaults();

CertKeyPair root_cert;

RETURN_NOT_OK(ExampleTlsCertificateRoot(&root_cert));

options.tls_root_certs = root_cert.pem_cert;

return FlightClient::Connect(location_, options).Value(&client_);

}

protected:

Location location_;

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,

std::shared_ptr<ServerMiddleware>* middleware) override {

return MakeFlightError(FlightStatusCode::Unauthenticated, "All calls are rejected");

}

public:

CountingServerMiddleware(std::atomic<int>* successful, std::atomic<int>* failed)

: successful_(successful), failed_(failed) {}

void SendingHeaders(AddCallHeaders* outgoing_headers) override {}

void CallCompleted(const Status& status) override {

if (status.ok()) {

ARROW_IGNORE_EXPR((*successful_)++);

} else {

ARROW_IGNORE_EXPR((*failed_)++);

}

}

std::string name() const override { return "CountingServerMiddleware"; }

private:

std::atomic<int>* successful_;

std::atomic<int>* failed_;

public:

CountingServerMiddlewareFactory() : successful_(0), failed_(0) {}

Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,

std::shared_ptr<ServerMiddleware>* middleware) override {

*middleware = std::make_shared<CountingServerMiddleware>(&successful_, &failed_);

return Status::OK();

}

std::atomic<int> successful_;

std::atomic<int> failed_;

public:

explicit TracingTestServerMiddleware(const std::string& current_span_id)

: span_id(current_span_id) {}

void SendingHeaders(AddCallHeaders* outgoing_headers) override {}

void CallCompleted(const Status& status) override {}

std::string name() const override { return "TracingTestServerMiddleware"; }

std::string span_id;

public:

TracingTestServerMiddlewareFactory() {}

Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,

std::shared_ptr<ServerMiddleware>* middleware) override {

const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =

incoming_headers.equal_range("x-tracing-span-id");

if (iter_pair.first != iter_pair.second) {

const std::string_view& value = (*iter_pair.first).second;

*middleware = std::make_shared<TracingTestServerMiddleware>(std::string(value));

}

return Status::OK();

}

const std::string& key,

const std::string& prefix) {

// Lambda function to compare characters without case sensitivity.

auto char_compare = [](const char& char1, const char& char2) {

return (::toupper(char1) == ::toupper(char2));

};

auto iter = incoming_headers.find(key);

if (iter == incoming_headers.end()) {

return "";

}

const std::string val(iter->second);

if (val.size() > prefix.length()) {

if (std::equal(val.begin(), val.begin() + prefix.length(), prefix.begin(),

char_compare)) {

return val.substr(prefix.length());

}

}

return "";

public:

void SendingHeaders(AddCallHeaders* outgoing_headers) override {

outgoing_headers->AddHeader(kAuthHeader, std::string(kBearerPrefix) + kBearerToken);

}

void CallCompleted(const Status& status) override {}

std::string name() const override { return "HeaderAuthServerMiddleware"; }

std::string& password) {

std::string encoded_credentials =

FindKeyValPrefixInCallHeaders(incoming_headers, kAuthHeader, kBasicPrefix);

std::stringstream decoded_stream(arrow::util::base64_decode(encoded_credentials));

std::getline(decoded_stream, username, ':');

std::getline(decoded_stream, password, ':');

public:

HeaderAuthServerMiddlewareFactory() {}

Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,

std::shared_ptr<ServerMiddleware>* middleware) override {

std::string username, password;

ParseBasicHeader(incoming_headers, username, password);

if ((username == kValidUsername) && (password == kValidPassword)) {

*middleware = std::make_shared<HeaderAuthServerMiddleware>();

} else if ((username == kInvalidUsername) && (password == kInvalidPassword)) {

return MakeFlightError(FlightStatusCode::Unauthenticated, "Invalid credentials");

}

return Status::OK();

}

public:

explicit BearerAuthServerMiddleware(const CallHeaders& incoming_headers, bool* isValid)

: isValid_(isValid) {

incoming_headers_ = incoming_headers;

}

void SendingHeaders(AddCallHeaders* outgoing_headers) override {

std::string bearer_token =

FindKeyValPrefixInCallHeaders(incoming_headers_, kAuthHeader, kBearerPrefix);

*isValid_ = (bearer_token == std::string(kBearerToken));

}

void CallCompleted(const Status& status) override {}

std::string name() const override { return "BearerAuthServerMiddleware"; }

private:

CallHeaders incoming_headers_;

bool* isValid_;

public:

BearerAuthServerMiddlewareFactory() : isValid_(false) {}

Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,

std::shared_ptr<ServerMiddleware>* middleware) override {

const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =

incoming_headers.equal_range(kAuthHeader);

if (iter_pair.first != iter_pair.second) {

*middleware =

std::make_shared<BearerAuthServerMiddleware>(incoming_headers, &isValid_);

}

return Status::OK();

}

bool GetIsValid() { return isValid_; }

private:

bool isValid_;

public:

explicit PropagatingClientMiddleware(std::atomic<int>* received_headers,

std::vector<Status>* recorded_status)

: received_headers_(received_headers), recorded_status_(recorded_status) {}

void SendingHeaders(AddCallHeaders* outgoing_headers) {

// Pick up the span ID from thread locals. We have to use a

// thread-local for communication, since we aren't even

// instantiated until after the application code has already

// started the call (and so there's no chance for application code

// to pass us parameters directly).

outgoing_headers->AddHeader("x-tracing-span-id", current_span_id);

}

void ReceivedHeaders(const CallHeaders& incoming_headers) { (*received_headers_)++; }

void CallCompleted(const Status& status) { recorded_status_->push_back(status); }

private:

std::atomic<int>* received_headers_;

std::vector<Status>* recorded_status_;

public:

void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {

recorded_calls_.push_back(info.method);

*middleware = std::make_unique<PropagatingClientMiddleware>(&received_headers_,

&recorded_status_);

}

void Reset() {

recorded_calls_.clear();

recorded_status_.clear();

received_headers_.fetch_and(0);

}

std::vector<FlightMethod> recorded_calls_;

std::vector<Status> recorded_status_;

std::atomic<int> received_headers_;

Status DoAction(const ServerCallContext& context, const Action& action,

std::unique_ptr<ResultStream>* result) override {

std::shared_ptr<Buffer> buf;

const ServerMiddleware* middleware = context.GetMiddleware("tracing");

if (middleware == nullptr || middleware->name() != "TracingTestServerMiddleware") {

buf = Buffer::FromString("");

} else {

buf = Buffer::FromString(((const TracingTestServerMiddleware*)middleware)->span_id);

}

*result = std::make_unique<SimpleResultStream>(std::vector<Result>{Result{buf}});

return Status::OK();

}

Status DoAction(const ServerCallContext& context, const Action& action,

std::unique_ptr<ResultStream>* result) override {

std::string msg = "error_message";

auto buf = Buffer::FromString("");

std::shared_ptr<FlightStatusDetail> flightStatusDetail(

new FlightStatusDetail(FlightStatusCode::Failed, msg));

*result = std::make_unique<SimpleResultStream>(std::vector<Result>{Result{buf}});

return Status(StatusCode::ExecutionError, "test failed", flightStatusDetail);

}

public:

explicit PropagatingTestServer(std::unique_ptr<FlightClient> client)

: client_(std::move(client)) {}

Status DoAction(const ServerCallContext& context, const Action& action,

std::unique_ptr<ResultStream>* result) override {

const ServerMiddleware* middleware = context.GetMiddleware("tracing");

if (middleware == nullptr || middleware->name() != "TracingTestServerMiddleware") {

current_span_id = "";

} else {

current_span_id = ((const TracingTestServerMiddleware*)middleware)->span_id;

}

return client_->DoAction(action).Value(result);

}

private:

std::unique_ptr<FlightClient> client_;

public:

void SetUp() {

ASSERT_OK(MakeServer<AppMetadataTestServer>(

&server_, &client_,

[](FlightServerOptions* options) {

options->middleware.push_back(

{"reject", std::make_shared<RejectServerMiddlewareFactory>()});

return Status::OK();

},

[](FlightClientOptions* options) { return Status::OK(); }));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

public:

void SetUp() {

request_counter_ = std::make_shared<CountingServerMiddlewareFactory>();

ASSERT_OK(MakeServer<AppMetadataTestServer>(

&server_, &client_,

[&](FlightServerOptions* options) {

options->middleware.push_back({"request_counter", request_counter_});

return Status::OK();

},

[](FlightClientOptions* options) { return Status::OK(); }));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::shared_ptr<CountingServerMiddlewareFactory> request_counter_;

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

public:

void SetUp() {

server_middleware_ = std::make_shared<TracingTestServerMiddlewareFactory>();

second_client_middleware_ = std::make_shared<PropagatingClientMiddlewareFactory>();

client_middleware_ = std::make_shared<PropagatingClientMiddlewareFactory>();

std::unique_ptr<FlightClient> server_client;

ASSERT_OK(MakeServer<ReportContextTestServer>(

&second_server_, &server_client,

[&](FlightServerOptions* options) {

options->middleware.push_back({"tracing", server_middleware_});

return Status::OK();

},

[&](FlightClientOptions* options) {

options->middleware.push_back(second_client_middleware_);

return Status::OK();

}));

ASSERT_OK(MakeServer<PropagatingTestServer>(

&first_server_, &client_,

[&](FlightServerOptions* options) {

options->middleware.push_back({"tracing", server_middleware_});

return Status::OK();

},

[&](FlightClientOptions* options) {

options->middleware.push_back(client_middleware_);

return Status::OK();

},

std::move(server_client)));

}

void ValidateStatus(const Status& status, const FlightMethod& method) {

ASSERT_EQ(1, client_middleware_->received_headers_);

ASSERT_EQ(method, client_middleware_->recorded_calls_.at(0));

ASSERT_EQ(status.code(), client_middleware_->recorded_status_.at(0).code());

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(first_server_->Shutdown());

ASSERT_OK(second_server_->Shutdown());

}

void CheckHeader(const std::string& header, const std::string& value,

const CallHeaders::const_iterator& it) {

// Construct a string_view before comparison to satisfy MSVC

std::string_view header_view(header.data(), header.length());

std::string_view value_view(value.data(), value.length());

ASSERT_EQ(header_view, (*it).first);

ASSERT_EQ(value_view, (*it).second);

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> first_server_;

std::unique_ptr<FlightServerBase> second_server_;

std::shared_ptr<TracingTestServerMiddlewareFactory> server_middleware_;

std::shared_ptr<PropagatingClientMiddlewareFactory> second_client_middleware_;

std::shared_ptr<PropagatingClientMiddlewareFactory> client_middleware_;

public:

void SetUp() {

ASSERT_OK(MakeServer<ErrorMiddlewareServer>(

&server_, &client_, [](FlightServerOptions* options) { return Status::OK(); },

[](FlightClientOptions* options) { return Status::OK(); }));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

public:

void SetUp() {

header_middleware_ = std::make_shared<HeaderAuthServerMiddlewareFactory>();

bearer_middleware_ = std::make_shared<BearerAuthServerMiddlewareFactory>();

std::pair<std::string, std::string> bearer = make_pair(

kAuthHeader, std::string(kBearerPrefix) + " " + std::string(kBearerToken));

ASSERT_OK(MakeServer<HeaderAuthTestServer>(

&server_, &client_,

[&](FlightServerOptions* options) {

options->auth_handler = std::make_unique<NoOpAuthHandler>();

options->middleware.push_back({"header-auth-server", header_middleware_});

options->middleware.push_back({"bearer-auth-server", bearer_middleware_});

return Status::OK();

},

[&](FlightClientOptions* options) { return Status::OK(); }));

}

void RunValidClientAuth() {

arrow::Result<std::pair<std::string, std::string>> bearer_result =

client_->AuthenticateBasicToken({}, kValidUsername, kValidPassword);

ASSERT_OK(bearer_result.status());

ASSERT_EQ(bearer_result.ValueOrDie().first, kAuthHeader);

ASSERT_EQ(bearer_result.ValueOrDie().second,

(std::string(kBearerPrefix) + kBearerToken));

std::unique_ptr<FlightListing> listing;

FlightCallOptions call_options;

call_options.headers.push_back(bearer_result.ValueOrDie());

ASSERT_OK_AND_ASSIGN(listing, client_->ListFlights(call_options, {}));

ASSERT_TRUE(bearer_middleware_->GetIsValid());

}

void RunInvalidClientAuth() {

arrow::Result<std::pair<std::string, std::string>> bearer_result =

client_->AuthenticateBasicToken({}, kInvalidUsername, kInvalidPassword);

ASSERT_RAISES(IOError, bearer_result.status());

ASSERT_THAT(bearer_result.status().message(),

::testing::HasSubstr("Invalid credentials"));

}

void TearDown() {

ASSERT_OK(client_->Close());

ASSERT_OK(server_->Shutdown());

}

protected:

std::unique_ptr<FlightClient> client_;

std::unique_ptr<FlightServerBase> server_;

std::shared_ptr<HeaderAuthServerMiddlewareFactory> header_middleware_;

std::shared_ptr<BearerAuthServerMiddlewareFactory> bearer_middleware_;

Action action;

// Run action1

action.type = "action1";

action.body = Buffer::FromString("action1-content");

ASSERT_OK_AND_ASSIGN(auto stream, client_->DoAction(action));

Status s = stream->Next().status();

ASSERT_FALSE(s.ok());

std::shared_ptr<FlightStatusDetail> flightStatusDetail =

FlightStatusDetail::UnwrapStatus(s);

ASSERT_TRUE(flightStatusDetail);

ASSERT_EQ(flightStatusDetail->extra_info(), "error_message");

ASSERT_OK_AND_ASSIGN(auto listing, client_->ListFlights());

ASSERT_TRUE(listing != nullptr);

std::vector<FlightInfo> flights = ExampleFlightInfo();

std::unique_ptr<FlightInfo> info;

for (const FlightInfo& flight : flights) {

ASSERT_OK_AND_ASSIGN(info, listing->Next());

AssertEqual(flight, *info);

}

ASSERT_OK_AND_ASSIGN(info, listing->Next());

ASSERT_TRUE(info == nullptr);

ASSERT_OK_AND_ASSIGN(info, listing->Next());

ASSERT_TRUE(info == nullptr);

ASSERT_OK_AND_ASSIGN(auto listing, client_->ListFlights(FlightCallOptions(), {"foo"}));

std::unique_ptr<FlightInfo> info;

ASSERT_OK_AND_ASSIGN(info, listing->Next());

ASSERT_TRUE(info == nullptr);

}

TEST_F(TestFlightClient, GetFlightInfo) {

auto descr = FlightDescriptor::Path({"examples", "ints"});

ASSERT_OK_AND_ASSIGN(auto info, client_->GetFlightInfo(descr));

ASSERT_NE(info, nullptr);

std::vector<FlightInfo> flights = ExampleFlightInfo();

AssertEqual(flights[0], *info);

}

TEST_F(TestFlightClient, GetSchema) {

auto descr = FlightDescriptor::Path({"examples", "ints"});

ipc::DictionaryMemo dict_memo;

ASSERT_OK_AND_ASSIGN(auto schema_result, client_->GetSchema(descr));

ASSERT_NE(schema_result, nullptr);

ASSERT_OK(schema_result->GetSchema(&dict_memo));

}

TEST_F(TestFlightClient, GetFlightInfoNotFound) {

auto descr = FlightDescriptor::Path({"examples", "things"});

// XXX Ideally should be Invalid (or KeyError), but gRPC doesn't support

// multiple error codes.

auto st = client_->GetFlightInfo(descr).status();

ASSERT_RAISES(Invalid, st);

ASSERT_NE(st.message().find("Flight not found"), std::string::npos);

}

TEST_F(TestFlightClient, ListActions) {

ASSERT_OK_AND_ASSIGN(std::vector<ActionType> actions, client_->ListActions());

std::vector<ActionType> expected = ExampleActionTypes();

EXPECT_THAT(actions, ::testing::ContainerEq(expected));

}

TEST_F(TestFlightClient, DoAction) {

Action action;

std::unique_ptr<Result> result;

// Run action1

action.type = "action1";

const std::string action1_value = "action1-content";

action.body = Buffer::FromString(action1_value);

ASSERT_OK_AND_ASSIGN(auto stream, client_->DoAction(action));

for (int i = 0; i < 3; ++i) {

ASSERT_OK_AND_ASSIGN(result, stream->Next());

std::string expected = action1_value + "-part" + std::to_string(i);

ASSERT_EQ(expected, result->body->ToString());

}

// stream consumed

ASSERT_OK_AND_ASSIGN(result, stream->Next());

ASSERT_EQ(nullptr, result);

// Run action2, no results

action.type = "action2";

ASSERT_OK_AND_ASSIGN(stream, client_->DoAction(action));

ASSERT_OK_AND_ASSIGN(result, stream->Next());

ASSERT_EQ(nullptr, result);

}

TEST_F(TestFlightClient, RoundTripStatus) {

const auto descr = FlightDescriptor::Command("status-outofmemory");

const auto status = client_->GetFlightInfo(descr).status();

ASSERT_RAISES(OutOfMemory, status);

}