1 /*
2 * Copyright 2022 Google LLC
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #include "fcp/client/http/protocol_request_helper.h"
17
18 #include "google/protobuf/any.pb.h"
19 #include "fcp/base/time_util.h"
20 #include "fcp/client/http/testing/test_helpers.h"
21 #include "fcp/client/test_helpers.h"
22 #include "fcp/protos/federatedcompute/secure_aggregations.pb.h"
23 #include "fcp/protos/federatedcompute/task_assignments.pb.h"
24 #include "fcp/testing/testing.h"
25
26 namespace fcp {
27 namespace client {
28 namespace http {
29 namespace {
30
31 using ::google::internal::federatedcompute::v1::AdvertiseKeysMetadata;
32 using ::google::internal::federatedcompute::v1::ForwardingInfo;
33 using ::google::internal::federatedcompute::v1::ShareKeysMetadata;
34 using ::google::internal::federatedcompute::v1::StartTaskAssignmentMetadata;
35 using ::google::internal::federatedcompute::v1::
36 SubmitSecureAggregationResultMetadata;
37 using ::google::longrunning::Operation;
38 using ::google::protobuf::Any;
39 using ::testing::_;
40 using ::testing::ContainerEq;
41 using ::testing::HasSubstr;
42 using ::testing::IsEmpty;
43 using ::testing::MockFunction;
44 using ::testing::NiceMock;
45 using ::testing::Return;
46 using ::testing::StrictMock;
47 using ::testing::UnorderedElementsAre;
48
49 constexpr absl::string_view kApiKey = "API_KEY";
50
51 class MockClock : public Clock {
52 public:
53 MOCK_METHOD(absl::Time, Now, (), (override));
54 MOCK_METHOD(void, Sleep, (absl::Duration duration), (override));
55
56 protected:
57 MOCK_METHOD(absl::Time, NowLocked, (), (override));
58 MOCK_METHOD(void, ScheduleWakeup, (absl::Time wakeup_time), (override));
59 };
60
VerifyInMemoryHttpResponse(const InMemoryHttpResponse & response,int code,absl::string_view content_encoding,absl::string_view body)61 void VerifyInMemoryHttpResponse(const InMemoryHttpResponse& response, int code,
62 absl::string_view content_encoding,
63 absl::string_view body) {
64 EXPECT_EQ(response.code, code);
65 EXPECT_EQ(response.content_encoding, content_encoding);
66 EXPECT_EQ(response.body, body);
67 }
68
CreatePendingOperation(const std::string operation_name)69 Operation CreatePendingOperation(const std::string operation_name) {
70 Operation operation;
71 operation.set_done(false);
72 operation.set_name(operation_name);
73 return operation;
74 }
75
CreatePendingOperation(const std::string operation_name,const Any & metadata)76 Operation CreatePendingOperation(const std::string operation_name,
77 const Any& metadata) {
78 Operation operation;
79 operation.set_done(false);
80 operation.set_name(operation_name);
81 *operation.mutable_metadata() = metadata;
82 return operation;
83 }
84
85 // Creates a 'done' `Operation`, with the given already-packed-into-`Any`
86 // result.
CreateDoneOperation(const Any & packed_inner_result)87 Operation CreateDoneOperation(const Any& packed_inner_result) {
88 Operation operation;
89 operation.set_done(true);
90 *operation.mutable_response() = packed_inner_result;
91 return operation;
92 }
93
CreateErrorOperation(const absl::StatusCode error_code,const std::string error_message)94 Operation CreateErrorOperation(const absl::StatusCode error_code,
95 const std::string error_message) {
96 Operation operation;
97 operation.set_done(true);
98 operation.mutable_error()->set_code(static_cast<int>(error_code));
99 operation.mutable_error()->set_message(error_message);
100 return operation;
101 }
102
TEST(ProtocolRequestCreatorTest,TestInvalidForwardingInfo)103 TEST(ProtocolRequestCreatorTest, TestInvalidForwardingInfo) {
104 // If a ForwardingInfo does not have a target_uri_prefix field set then the
105 // ProcessForwardingInfo call should fail.
106 ForwardingInfo forwarding_info;
107 EXPECT_THAT(ProtocolRequestCreator::Create(kApiKey, forwarding_info,
108 /*use_compression=*/false),
109 IsCode(INVALID_ARGUMENT));
110
111 (*forwarding_info.mutable_extra_request_headers())["x-header1"] =
112 "header-value1";
113 EXPECT_THAT(ProtocolRequestCreator::Create(kApiKey, forwarding_info,
114 /*use_compression=*/false),
115 IsCode(INVALID_ARGUMENT));
116 }
117
TEST(ProtocolRequestCreatorTest,CreateProtocolRequestInvalidSuffix)118 TEST(ProtocolRequestCreatorTest, CreateProtocolRequestInvalidSuffix) {
119 ProtocolRequestCreator creator("https://initial.uri", kApiKey, HeaderList(),
120 /*use_compression=*/false);
121 std::string uri_suffix = "v1/request";
122 ASSERT_THAT(
123 creator.CreateProtocolRequest(uri_suffix, QueryParams(),
124 HttpRequest::Method::kPost, "request_body",
125 /*is_protobuf_encoded=*/false),
126 IsCode(absl::StatusCode::kInvalidArgument));
127 }
128
TEST(ProtocolRequestCreatorTest,CreateProtocolRequest)129 TEST(ProtocolRequestCreatorTest, CreateProtocolRequest) {
130 ProtocolRequestCreator creator("https://initial.uri", kApiKey, HeaderList(),
131 /*use_compression=*/false);
132 std::string expected_body = "expected_body";
133 auto request = creator.CreateProtocolRequest(
134 "/v1/request", QueryParams(), HttpRequest::Method::kPost, expected_body,
135 /*is_protobuf_encoded=*/false);
136
137 ASSERT_OK(request);
138 EXPECT_EQ((*request)->uri(), "https://initial.uri/v1/request");
139 EXPECT_EQ((*request)->method(), HttpRequest::Method::kPost);
140 EXPECT_THAT(
141 (*request)->extra_headers(),
142 UnorderedElementsAre(
143 Header{"x-goog-api-key", "API_KEY"},
144 Header{"Content-Length", std::to_string(expected_body.size())}));
145 EXPECT_TRUE((*request)->HasBody());
146 std::string actual_body;
147 actual_body.resize(expected_body.size());
148 ASSERT_OK((*request)->ReadBody(actual_body.data(), expected_body.size()));
149 EXPECT_EQ(actual_body, expected_body);
150 }
151
TEST(ProtocolRequestCreatorTest,CreateProtobufEncodedProtocolRequest)152 TEST(ProtocolRequestCreatorTest, CreateProtobufEncodedProtocolRequest) {
153 ProtocolRequestCreator creator("https://initial.uri", kApiKey, HeaderList(),
154 /*use_compression=*/false);
155 std::string expected_body = "expected_body";
156 auto request = creator.CreateProtocolRequest(
157 "/v1/request", QueryParams(), HttpRequest::Method::kPost, expected_body,
158 /*is_protobuf_encoded=*/true);
159
160 ASSERT_OK(request);
161 EXPECT_EQ((*request)->uri(), "https://initial.uri/v1/request?%24alt=proto");
162 EXPECT_EQ((*request)->method(), HttpRequest::Method::kPost);
163 EXPECT_THAT((*request)->extra_headers(),
164 UnorderedElementsAre(
165 Header{"x-goog-api-key", "API_KEY"},
166 Header{"Content-Length", absl::StrCat(expected_body.size())},
167 Header{"Content-Type", "application/x-protobuf"}));
168 EXPECT_TRUE((*request)->HasBody());
169 std::string actual_body;
170 actual_body.resize(expected_body.size());
171 ASSERT_OK((*request)->ReadBody(actual_body.data(), actual_body.size()));
172 EXPECT_EQ(actual_body, expected_body);
173 }
174
TEST(ProtocolRequestCreatorTest,CreateGetOperationRequest)175 TEST(ProtocolRequestCreatorTest, CreateGetOperationRequest) {
176 ProtocolRequestCreator creator("https://initial.uri", kApiKey, HeaderList(),
177 /*use_compression=*/false);
178 std::string operation_name = "my_operation";
179 auto request = creator.CreateGetOperationRequest(operation_name);
180 ASSERT_OK(request);
181 EXPECT_EQ((*request)->uri(),
182 "https://initial.uri/v1/my_operation?%24alt=proto");
183 EXPECT_EQ((*request)->method(), HttpRequest::Method::kGet);
184 EXPECT_THAT((*request)->extra_headers(),
185 UnorderedElementsAre(Header{"x-goog-api-key", "API_KEY"}));
186 EXPECT_FALSE((*request)->HasBody());
187 }
188
TEST(ProtocolRequestCreatorTest,CreateCancelOperationRequest)189 TEST(ProtocolRequestCreatorTest, CreateCancelOperationRequest) {
190 ProtocolRequestCreator creator("https://initial.uri", kApiKey, HeaderList(),
191 /*use_compression=*/false);
192 std::string operation_name = "my_operation";
193 auto request = creator.CreateCancelOperationRequest(operation_name);
194 ASSERT_OK(request);
195 EXPECT_EQ((*request)->uri(),
196 "https://initial.uri/v1/my_operation:cancel?%24alt=proto");
197 EXPECT_EQ((*request)->method(), HttpRequest::Method::kGet);
198 EXPECT_THAT((*request)->extra_headers(),
199 UnorderedElementsAre(Header{"x-goog-api-key", "API_KEY"}));
200 EXPECT_FALSE((*request)->HasBody());
201 }
202
203 class ProtocolRequestHelperTest : public ::testing::Test {
204 public:
ProtocolRequestHelperTest()205 ProtocolRequestHelperTest()
206 : interruptible_runner_(
207 &mock_log_manager_, mock_should_abort_.AsStdFunction(),
208 InterruptibleRunner::TimingConfig{
209 .polling_period = absl::ZeroDuration(),
210 .graceful_shutdown_period = absl::InfiniteDuration(),
211 .extended_shutdown_period = absl::InfiniteDuration()},
212 InterruptibleRunner::DiagnosticsConfig{
213 .interrupted = ProdDiagCode::BACKGROUND_TRAINING_INTERRUPT_HTTP,
214 .interrupt_timeout =
215 ProdDiagCode::BACKGROUND_TRAINING_INTERRUPT_HTTP_TIMED_OUT,
216 .interrupted_extended = ProdDiagCode::
217 BACKGROUND_TRAINING_INTERRUPT_HTTP_EXTENDED_COMPLETED,
218 .interrupt_timeout_extended = ProdDiagCode::
219 BACKGROUND_TRAINING_INTERRUPT_HTTP_EXTENDED_TIMED_OUT}),
220 initial_request_creator_("https://initial.uri", kApiKey, HeaderList(),
221 /*use_compression=*/false),
222 protocol_request_helper_(&mock_http_client_, &bytes_downloaded_,
223 &bytes_uploaded_, network_stopwatch_.get(),
224 &mock_clock_) {}
225
226 protected:
TearDown()227 void TearDown() override {
228 // Regardless of the outcome of the test (or the protocol interaction being
229 // tested), network usage must always be reflected in the network stats.
230 HttpRequestHandle::SentReceivedBytes sent_received_bytes =
231 mock_http_client_.TotalSentReceivedBytes();
232 EXPECT_THAT(bytes_downloaded_, sent_received_bytes.received_bytes);
233 EXPECT_THAT(bytes_uploaded_, sent_received_bytes.sent_bytes);
234 }
235
236 StrictMock<MockClock> mock_clock_;
237 StrictMock<MockHttpClient> mock_http_client_;
238
239 NiceMock<MockLogManager> mock_log_manager_;
240 NiceMock<MockFunction<bool()>> mock_should_abort_;
241
242 int64_t bytes_downloaded_ = 0;
243 int64_t bytes_uploaded_ = 0;
244 std::unique_ptr<WallClockStopwatch> network_stopwatch_ =
245 WallClockStopwatch::Create();
246
247 InterruptibleRunner interruptible_runner_;
248 ProtocolRequestCreator initial_request_creator_;
249 // The class under test.
250 ProtocolRequestHelper protocol_request_helper_;
251 };
252
GetFakeAnyProto()253 Any GetFakeAnyProto() {
254 Any fake_any;
255 fake_any.set_type_url("the_type_url");
256 *fake_any.mutable_value() = "the_value";
257 return fake_any;
258 }
259
TEST_F(ProtocolRequestHelperTest,TestForwardingInfoIsPassedAlongCorrectly)260 TEST_F(ProtocolRequestHelperTest, TestForwardingInfoIsPassedAlongCorrectly) {
261 // The initial request should use the initial entry point URI and an empty set
262 // of headers.
263 EXPECT_CALL(mock_http_client_,
264 PerformSingleRequest(SimpleHttpRequestMatcher(
265 "https://initial.uri/suffix1", HttpRequest::Method::kPost,
266 // This request has a response body, so the HttpClient will
267 // add this header automatically.
268 ContainerEq(HeaderList{{"x-goog-api-key", "API_KEY"},
269 {"Content-Length", "5"}}),
270 "body1")))
271 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response1")));
272 auto request_creator = std::make_unique<ProtocolRequestCreator>(
273 "https://initial.uri", kApiKey, HeaderList(),
274 /*use_compression=*/false);
275 auto http_request = request_creator->CreateProtocolRequest(
276 "/suffix1", QueryParams(), HttpRequest::Method::kPost, "body1",
277 /*is_protobuf_encoded=*/false);
278 ASSERT_OK(http_request);
279 auto result = protocol_request_helper_.PerformProtocolRequest(
280 *std::move(http_request), interruptible_runner_);
281 ASSERT_OK(result);
282 VerifyInMemoryHttpResponse(*result, 200, "", "response1");
283
284 {
285 // Process some fake ForwardingInfo.
286 ForwardingInfo forwarding_info1;
287 forwarding_info1.set_target_uri_prefix("https://second.uri/");
288 (*forwarding_info1.mutable_extra_request_headers())["x-header1"] =
289 "header-value1";
290 (*forwarding_info1.mutable_extra_request_headers())["x-header2"] =
291 "header-value2";
292 auto new_request_creator = ProtocolRequestCreator::Create(
293 kApiKey, forwarding_info1, /*use_compression=*/false);
294 ASSERT_OK(new_request_creator);
295 request_creator = std::move(*new_request_creator);
296 }
297
298 // The next series of requests should now use the ForwardingInfo (incl. use
299 // the "https://second.uri/" prefix, and include the headers). Note that we
300 // must use UnorderedElementsAre since the iteration order of the headers in
301 // the `ForwardingInfo` is undefined.
302 EXPECT_CALL(mock_http_client_,
303 PerformSingleRequest(SimpleHttpRequestMatcher(
304 "https://second.uri/suffix2", HttpRequest::Method::kGet,
305 UnorderedElementsAre(Header{"x-goog-api-key", "API_KEY"},
306 Header{"x-header1", "header-value1"},
307 Header{"x-header2", "header-value2"}),
308 "")))
309 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response2")));
310 http_request = request_creator->CreateProtocolRequest(
311 "/suffix2", QueryParams(), HttpRequest::Method::kGet, "",
312 /*is_protobuf_encoded=*/false);
313 ASSERT_OK(http_request);
314 result = protocol_request_helper_.PerformProtocolRequest(
315 *std::move(http_request), interruptible_runner_);
316 ASSERT_OK(result);
317 EXPECT_EQ(result->body, "response2");
318
319 EXPECT_CALL(mock_http_client_,
320 PerformSingleRequest(SimpleHttpRequestMatcher(
321 "https://second.uri/suffix3", HttpRequest::Method::kPut,
322 UnorderedElementsAre(Header{"x-goog-api-key", "API_KEY"},
323 Header{"x-header1", "header-value1"},
324 Header{"x-header2", "header-value2"},
325 // This request has a response body, so
326 // the HttpClient will add this header
327 // automatically.
328 Header{"Content-Length", "5"}),
329 "body3")))
330 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response3")));
331 http_request = request_creator->CreateProtocolRequest(
332 "/suffix3", QueryParams(), HttpRequest::Method::kPut, "body3",
333 /*is_protobuf_encoded=*/false);
334 ASSERT_OK(http_request);
335 result = protocol_request_helper_.PerformProtocolRequest(
336 *std::move(http_request), interruptible_runner_);
337 ASSERT_OK(result);
338 EXPECT_EQ(result->body, "response3");
339
340 {
341 // Process some more fake ForwardingInfo (without any headers this time).
342 ForwardingInfo forwarding_info2;
343 forwarding_info2.set_target_uri_prefix("https://third.uri");
344 auto new_request_creator = ProtocolRequestCreator::Create(
345 kApiKey, forwarding_info2, /*use_compression=*/false);
346 ASSERT_OK(new_request_creator);
347 request_creator = std::move(*new_request_creator);
348 }
349
350 // The next request should now use the latest ForwardingInfo again (i.e. use
351 // the "https://third.uri/" prefix, and not specify any headers).
352 EXPECT_CALL(mock_http_client_,
353 PerformSingleRequest(SimpleHttpRequestMatcher(
354 "https://third.uri/suffix4", HttpRequest::Method::kPost,
355 // This request has a response body, so the HttpClient will
356 // add this header automatically.
357 ContainerEq(HeaderList{
358 {"x-goog-api-key", "API_KEY"},
359 {"Content-Length", "5"},
360 }),
361 "body4")))
362 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response4")));
363 http_request = request_creator->CreateProtocolRequest(
364 "/suffix4", QueryParams(), HttpRequest::Method::kPost, "body4",
365 /*is_protobuf_encoded=*/false);
366 ASSERT_OK(http_request);
367 result = protocol_request_helper_.PerformProtocolRequest(
368 *std::move(http_request), interruptible_runner_);
369 ASSERT_OK(result);
370 EXPECT_EQ(result->body, "response4");
371 }
372
TEST_F(ProtocolRequestHelperTest,TestPollOperationInvalidOperationName)373 TEST_F(ProtocolRequestHelperTest, TestPollOperationInvalidOperationName) {
374 absl::StatusOr<Operation> result =
375 protocol_request_helper_.PollOperationResponseUntilDone(
376 CreatePendingOperation("invalid_operation_name"),
377 initial_request_creator_, interruptible_runner_);
378 EXPECT_THAT(result.status(), IsCode(INVALID_ARGUMENT));
379 EXPECT_THAT(result.status().message(), HasSubstr("invalid name"));
380 }
381
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseImmediateSuccess)382 TEST_F(ProtocolRequestHelperTest, TestPollOperationResponseImmediateSuccess) {
383 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
384 absl::StatusOr<Operation> result =
385 protocol_request_helper_.PollOperationResponseUntilDone(
386 expected_response, initial_request_creator_, interruptible_runner_);
387 ASSERT_OK(result);
388 EXPECT_THAT(*result, EqualsProto(expected_response));
389 }
390
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseImmediateOperationError)391 TEST_F(ProtocolRequestHelperTest,
392 TestPollOperationResponseImmediateOperationError) {
393 Operation expected_response =
394 CreateErrorOperation(ALREADY_EXISTS, "some error");
395 absl::StatusOr<Operation> result =
396 protocol_request_helper_.PollOperationResponseUntilDone(
397 expected_response, initial_request_creator_, interruptible_runner_);
398 ASSERT_OK(result);
399 EXPECT_THAT(*result, EqualsProto(expected_response));
400 }
401
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseSuccessAfterPolling)402 TEST_F(ProtocolRequestHelperTest,
403 TestPollOperationResponseSuccessAfterPolling) {
404 // Make the initial request return a pending Operation result. Note that we
405 // use a '#' character in the operation name to allow us to verify that it
406 // is properly URL-encoded.
407 Operation pending_operation_response =
408 CreatePendingOperation("operations/foo#bar");
409
410 // Then, after letting the operation get polled twice more, eventually
411 // return a fake response.
412 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
413 EXPECT_CALL(mock_http_client_,
414 PerformSingleRequest(SimpleHttpRequestMatcher(
415 // Note that the '#' character is encoded as "%23".
416 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
417 HttpRequest::Method::kGet, _, IsEmpty())))
418 .WillOnce(Return(FakeHttpResponse(
419 200, HeaderList(), pending_operation_response.SerializeAsString())))
420 .WillOnce(Return(FakeHttpResponse(
421 200, HeaderList(), pending_operation_response.SerializeAsString())))
422 .WillOnce(Return(FakeHttpResponse(
423 200, HeaderList(), expected_response.SerializeAsString())));
424 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500))).Times(3);
425 absl::StatusOr<Operation> result =
426 protocol_request_helper_.PollOperationResponseUntilDone(
427 pending_operation_response, initial_request_creator_,
428 interruptible_runner_);
429 ASSERT_OK(result);
430 EXPECT_THAT(*result, EqualsProto(expected_response));
431 }
432
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseErrorAfterPolling)433 TEST_F(ProtocolRequestHelperTest, TestPollOperationResponseErrorAfterPolling) {
434 // Make the initial request return a pending Operation result.
435 Operation pending_operation_response =
436 CreatePendingOperation("operations/foo#bar");
437
438 // Then, after letting the operation get polled twice more, eventually
439 // return a fake error.
440 Operation expected_response =
441 CreateErrorOperation(ALREADY_EXISTS, "some error");
442 EXPECT_CALL(mock_http_client_,
443 PerformSingleRequest(SimpleHttpRequestMatcher(
444 // Note that the '#' character is encoded as "%23".
445 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
446 HttpRequest::Method::kGet, _, IsEmpty())))
447 .WillOnce(Return(FakeHttpResponse(
448 200, HeaderList(), pending_operation_response.SerializeAsString())))
449 .WillOnce(Return(FakeHttpResponse(
450 200, HeaderList(), pending_operation_response.SerializeAsString())))
451 .WillOnce(Return(FakeHttpResponse(
452 200, HeaderList(), expected_response.SerializeAsString())));
453 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500))).Times(3);
454
455 absl::StatusOr<Operation> result =
456 protocol_request_helper_.PollOperationResponseUntilDone(
457 pending_operation_response, initial_request_creator_,
458 interruptible_runner_);
459 ASSERT_OK(result);
460 EXPECT_THAT(*result, EqualsProto(expected_response));
461 }
462
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseDifferentPollingIntervals)463 TEST_F(ProtocolRequestHelperTest,
464 TestPollOperationResponseDifferentPollingIntervals) {
465 StartTaskAssignmentMetadata metadata;
466 *metadata.mutable_polling_interval() =
467 TimeUtil::ConvertAbslToProtoDuration(absl::Milliseconds(2));
468 Any packed_metadata;
469 ASSERT_TRUE(packed_metadata.PackFrom(metadata));
470 StartTaskAssignmentMetadata metadata_2;
471 *metadata_2.mutable_polling_interval() =
472 TimeUtil::ConvertAbslToProtoDuration(absl::Milliseconds(3));
473 Any packed_metadata_2;
474 ASSERT_TRUE(packed_metadata_2.PackFrom(metadata_2));
475
476 // Make the initial request return a pending Operation result. Note that we
477 // use a '#' character in the operation name to allow us to verify that it
478 // is properly URL-encoded.
479 Operation pending_operation_response =
480 CreatePendingOperation("operations/foo#bar");
481
482 // Then, after letting the operation get polled twice more, eventually
483 // return a fake response.
484 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
485
486 EXPECT_CALL(mock_http_client_,
487 PerformSingleRequest(SimpleHttpRequestMatcher(
488 // Note that the '#' character is encoded as "%23".
489 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
490 HttpRequest::Method::kGet, _, IsEmpty())))
491 .WillOnce(Return(FakeHttpResponse(
492 200, HeaderList(),
493 CreatePendingOperation("operations/foo#bar", packed_metadata)
494 .SerializeAsString())))
495 .WillOnce(Return(FakeHttpResponse(
496 200, HeaderList(),
497 CreatePendingOperation("operations/foo#bar", packed_metadata_2)
498 .SerializeAsString())))
499 .WillOnce(Return(FakeHttpResponse(
500 200, HeaderList(), expected_response.SerializeAsString())));
501 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500)));
502 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(2)));
503 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(3)));
504 absl::StatusOr<Operation> result =
505 protocol_request_helper_.PollOperationResponseUntilDone(
506 pending_operation_response, initial_request_creator_,
507 interruptible_runner_);
508 ASSERT_OK(result);
509 EXPECT_THAT(*result, EqualsProto(expected_response));
510 }
511
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponsePollingIntervalTooHigh)512 TEST_F(ProtocolRequestHelperTest,
513 TestPollOperationResponsePollingIntervalTooHigh) {
514 StartTaskAssignmentMetadata metadata;
515 *metadata.mutable_polling_interval() =
516 TimeUtil::ConvertAbslToProtoDuration(absl::Hours(1));
517 Any packed_metadata;
518 ASSERT_TRUE(packed_metadata.PackFrom(metadata));
519
520 // Make the initial request return a pending Operation result. Note that we
521 // use a '#' character in the operation name to allow us to verify that it
522 // is properly URL-encoded.
523 Operation pending_operation_response =
524 CreatePendingOperation("operations/foo#bar");
525
526 // Then, after letting the operation get polled twice more, eventually
527 // return a fake response.
528 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
529
530 EXPECT_CALL(mock_http_client_,
531 PerformSingleRequest(SimpleHttpRequestMatcher(
532 // Note that the '#' character is encoded as "%23".
533 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
534 HttpRequest::Method::kGet, _, IsEmpty())))
535 .WillOnce(Return(FakeHttpResponse(
536 200, HeaderList(),
537 CreatePendingOperation("operations/foo#bar", packed_metadata)
538 .SerializeAsString())))
539 .WillOnce(Return(FakeHttpResponse(
540 200, HeaderList(), expected_response.SerializeAsString())));
541 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500)));
542 EXPECT_CALL(mock_clock_, Sleep(absl::Minutes(1)));
543 absl::StatusOr<Operation> result =
544 protocol_request_helper_.PollOperationResponseUntilDone(
545 pending_operation_response, initial_request_creator_,
546 interruptible_runner_);
547 ASSERT_OK(result);
548 EXPECT_THAT(*result, EqualsProto(expected_response));
549 }
550
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseAdvertiseKeysMetadata)551 TEST_F(ProtocolRequestHelperTest,
552 TestPollOperationResponseAdvertiseKeysMetadata) {
553 AdvertiseKeysMetadata metadata;
554 *metadata.mutable_polling_interval() =
555 TimeUtil::ConvertAbslToProtoDuration(absl::Milliseconds(2));
556 Any packed_metadata;
557 ASSERT_TRUE(packed_metadata.PackFrom(metadata));
558
559 // Make the initial request return a pending Operation result. Note that we
560 // use a '#' character in the operation name to allow us to verify that it
561 // is properly URL-encoded.
562 Operation pending_operation_response =
563 CreatePendingOperation("operations/foo#bar");
564
565 // Then, after letting the operation get polled twice more, eventually
566 // return a fake response.
567 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
568
569 EXPECT_CALL(mock_http_client_,
570 PerformSingleRequest(SimpleHttpRequestMatcher(
571 // Note that the '#' character is encoded as "%23".
572 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
573 HttpRequest::Method::kGet, _, IsEmpty())))
574 .WillOnce(Return(FakeHttpResponse(
575 200, HeaderList(),
576 CreatePendingOperation("operations/foo#bar", packed_metadata)
577 .SerializeAsString())))
578 .WillOnce(Return(FakeHttpResponse(
579 200, HeaderList(), expected_response.SerializeAsString())));
580 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500)));
581 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(2)));
582 absl::StatusOr<Operation> result =
583 protocol_request_helper_.PollOperationResponseUntilDone(
584 pending_operation_response, initial_request_creator_,
585 interruptible_runner_);
586 ASSERT_OK(result);
587 EXPECT_THAT(*result, EqualsProto(expected_response));
588 }
589
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseShareKeysMetadata)590 TEST_F(ProtocolRequestHelperTest, TestPollOperationResponseShareKeysMetadata) {
591 ShareKeysMetadata metadata;
592 *metadata.mutable_polling_interval() =
593 TimeUtil::ConvertAbslToProtoDuration(absl::Milliseconds(2));
594 Any packed_metadata;
595 ASSERT_TRUE(packed_metadata.PackFrom(metadata));
596
597 // Make the initial request return a pending Operation result. Note that we
598 // use a '#' character in the operation name to allow us to verify that it
599 // is properly URL-encoded.
600 Operation pending_operation_response =
601 CreatePendingOperation("operations/foo#bar");
602
603 // Then, after letting the operation get polled twice more, eventually
604 // return a fake response.
605 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
606
607 EXPECT_CALL(mock_http_client_,
608 PerformSingleRequest(SimpleHttpRequestMatcher(
609 // Note that the '#' character is encoded as "%23".
610 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
611 HttpRequest::Method::kGet, _, IsEmpty())))
612 .WillOnce(Return(FakeHttpResponse(
613 200, HeaderList(),
614 CreatePendingOperation("operations/foo#bar", packed_metadata)
615 .SerializeAsString())))
616 .WillOnce(Return(FakeHttpResponse(
617 200, HeaderList(), expected_response.SerializeAsString())));
618 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500)));
619 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(2)));
620 absl::StatusOr<Operation> result =
621 protocol_request_helper_.PollOperationResponseUntilDone(
622 pending_operation_response, initial_request_creator_,
623 interruptible_runner_);
624 ASSERT_OK(result);
625 EXPECT_THAT(*result, EqualsProto(expected_response));
626 }
627
TEST_F(ProtocolRequestHelperTest,TestPollOperationResponseSubmitSecureAggregationResultMetadata)628 TEST_F(ProtocolRequestHelperTest,
629 TestPollOperationResponseSubmitSecureAggregationResultMetadata) {
630 SubmitSecureAggregationResultMetadata metadata;
631 *metadata.mutable_polling_interval() =
632 TimeUtil::ConvertAbslToProtoDuration(absl::Milliseconds(2));
633 Any packed_metadata;
634 ASSERT_TRUE(packed_metadata.PackFrom(metadata));
635
636 // Make the initial request return a pending Operation result. Note that we
637 // use a '#' character in the operation name to allow us to verify that it
638 // is properly URL-encoded.
639 Operation pending_operation_response =
640 CreatePendingOperation("operations/foo#bar");
641
642 // Then, after letting the operation get polled twice more, eventually
643 // return a fake response.
644 Operation expected_response = CreateDoneOperation(GetFakeAnyProto());
645
646 EXPECT_CALL(mock_http_client_,
647 PerformSingleRequest(SimpleHttpRequestMatcher(
648 // Note that the '#' character is encoded as "%23".
649 "https://initial.uri/v1/operations/foo%23bar?%24alt=proto",
650 HttpRequest::Method::kGet, _, IsEmpty())))
651 .WillOnce(Return(FakeHttpResponse(
652 200, HeaderList(),
653 CreatePendingOperation("operations/foo#bar", packed_metadata)
654 .SerializeAsString())))
655 .WillOnce(Return(FakeHttpResponse(
656 200, HeaderList(), expected_response.SerializeAsString())));
657 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(500)));
658 EXPECT_CALL(mock_clock_, Sleep(absl::Milliseconds(2)));
659 absl::StatusOr<Operation> result =
660 protocol_request_helper_.PollOperationResponseUntilDone(
661 pending_operation_response, initial_request_creator_,
662 interruptible_runner_);
663 ASSERT_OK(result);
664 EXPECT_THAT(*result, EqualsProto(expected_response));
665 }
666
TEST_F(ProtocolRequestHelperTest,PerformMultipleRequestsSuccess)667 TEST_F(ProtocolRequestHelperTest, PerformMultipleRequestsSuccess) {
668 auto request_a = initial_request_creator_.CreateProtocolRequest(
669 "/v1/request_a", QueryParams(), HttpRequest::Method::kPost, "body1",
670 /*is_protobuf_encoded=*/false);
671 ASSERT_OK(request_a);
672 auto request_b = initial_request_creator_.CreateProtocolRequest(
673 "/v1/request_b", QueryParams(), HttpRequest::Method::kPost, "body2",
674 /*is_protobuf_encoded=*/false);
675 ASSERT_OK(request_b);
676 EXPECT_CALL(
677 mock_http_client_,
678 PerformSingleRequest(SimpleHttpRequestMatcher(
679 "https://initial.uri/v1/request_a", HttpRequest::Method::kPost,
680 // This request has a response body, so the HttpClient will
681 // add this header automatically.
682 ContainerEq(HeaderList{{"x-goog-api-key", "API_KEY"},
683 {"Content-Length", "5"}}),
684 "body1")))
685 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response1")));
686 EXPECT_CALL(
687 mock_http_client_,
688 PerformSingleRequest(SimpleHttpRequestMatcher(
689 "https://initial.uri/v1/request_b", HttpRequest::Method::kPost,
690 // This request has a response body, so the HttpClient will
691 // add this header automatically.
692 ContainerEq(HeaderList{{"x-goog-api-key", "API_KEY"},
693 {"Content-Length", "5"}}),
694 "body2")))
695 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response2")));
696
697 std::vector<std::unique_ptr<HttpRequest>> requests;
698 requests.push_back(std::move(*request_a));
699 requests.push_back(std::move(*request_b));
700
701 auto result = protocol_request_helper_.PerformMultipleProtocolRequests(
702 std::move(requests), interruptible_runner_);
703 ASSERT_OK(result);
704 auto response_1 = (*result)[0];
705 ASSERT_OK(response_1);
706 VerifyInMemoryHttpResponse(*response_1, 200, "", "response1");
707 auto response_2 = (*result)[1];
708 ASSERT_OK(response_2);
709 VerifyInMemoryHttpResponse(*response_2, 200, "", "response2");
710 }
711
TEST_F(ProtocolRequestHelperTest,PerformMultipleRequestsPartialFail)712 TEST_F(ProtocolRequestHelperTest, PerformMultipleRequestsPartialFail) {
713 std::string uri_suffix_a = "/v1/request_a";
714 auto request_a = initial_request_creator_.CreateProtocolRequest(
715 uri_suffix_a, QueryParams(), HttpRequest::Method::kPost, "body1",
716 /*is_protobuf_encoded=*/false);
717 ASSERT_OK(request_a);
718 std::string uri_suffix_b = "/v1/request_b";
719 auto request_b = initial_request_creator_.CreateProtocolRequest(
720 uri_suffix_b, QueryParams(), HttpRequest::Method::kPost, "body2",
721 /*is_protobuf_encoded=*/false);
722 ASSERT_OK(request_b);
723 EXPECT_CALL(
724 mock_http_client_,
725 PerformSingleRequest(SimpleHttpRequestMatcher(
726 "https://initial.uri/v1/request_a", HttpRequest::Method::kPost,
727 // This request has a response body, so the HttpClient will
728 // add this header automatically.
729 ContainerEq(HeaderList{{"x-goog-api-key", "API_KEY"},
730 {"Content-Length", "5"}}),
731 "body1")))
732 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "response1")));
733 EXPECT_CALL(
734 mock_http_client_,
735 PerformSingleRequest(SimpleHttpRequestMatcher(
736 "https://initial.uri/v1/request_b", HttpRequest::Method::kPost,
737 // This request has a response body, so the HttpClient will
738 // add this header automatically.
739 ContainerEq(HeaderList{{"x-goog-api-key", "API_KEY"},
740 {"Content-Length", "5"}}),
741 "body2")))
742 .WillOnce(
743 Return(FakeHttpResponse(404, HeaderList(), "failure_response")));
744
745 std::vector<std::unique_ptr<HttpRequest>> requests;
746 requests.push_back(std::move(*request_a));
747 requests.push_back(std::move(*request_b));
748
749 auto result = protocol_request_helper_.PerformMultipleProtocolRequests(
750 std::move(requests), interruptible_runner_);
751
752 ASSERT_OK(result);
753 auto response_1 = (*result)[0];
754 ASSERT_OK(response_1);
755 VerifyInMemoryHttpResponse(*response_1, 200, "", "response1");
756 auto response_2 = (*result)[1];
757 ASSERT_THAT(response_2, IsCode(absl::StatusCode::kNotFound));
758 }
759 } // anonymous namespace
760 } // namespace http
761 } // namespace client
762 } // namespace fcp
763