1 /* 2 * Copyright 2021 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 17 #ifndef FCP_CLIENT_FAKE_SERVER_H_ 18 #define FCP_CLIENT_FAKE_SERVER_H_ 19 20 #include <cstddef> 21 #include <string> 22 #include <tuple> 23 24 #include "grpcpp/impl/codegen/status.h" 25 #include "absl/status/status.h" 26 #include "absl/synchronization/notification.h" 27 #include "fcp/base/monitoring.h" 28 #include "fcp/client/grpc_bidi_stream.h" 29 #include "fcp/protocol/grpc_chunked_bidi_stream.h" 30 #include "fcp/protos/federated_api.grpc.pb.h" 31 #include "fcp/protos/federated_api.pb.h" 32 #include "grpcpp/impl/codegen/server_context.h" 33 34 namespace fcp { 35 namespace client { 36 namespace test { 37 38 class FakeServer 39 : public google::internal::federatedml::v2::FederatedTrainingApi::Service { 40 public: FakeServer()41 FakeServer() 42 : chunk_size_for_upload_(8192), 43 max_pending_chunks_(2), 44 compression_level_(google::internal::federatedml::v2::CompressionLevel:: 45 ZLIB_BEST_COMPRESSION) {} FakeServer(int32_t chunk_size_for_upload,int32_t max_pending_chunks,google::internal::federatedml::v2::CompressionLevel compression_level)46 FakeServer( 47 int32_t chunk_size_for_upload, int32_t max_pending_chunks, 48 google::internal::federatedml::v2::CompressionLevel compression_level) 49 : chunk_size_for_upload_(chunk_size_for_upload), 50 max_pending_chunks_(max_pending_chunks), 51 compression_level_(compression_level) {} 52 53 // FakeServer is neither copyable nor movable. 54 FakeServer(const FakeServer&) = delete; 55 FakeServer& operator=(const FakeServer&) = delete; 56 57 grpc::Status Session( 58 grpc::ServerContext* context, 59 grpc::ServerReaderWriter< 60 google::internal::federatedml::v2::ServerStreamMessage, 61 google::internal::federatedml::v2::ClientStreamMessage>* stream) 62 override; 63 void WaitForSessionDone(); 64 65 virtual absl::Status Handle( 66 const google::internal::federatedml::v2::ClientStreamMessage& request, 67 google::internal::federatedml::v2::ServerStreamMessage* first_reply, 68 ::fcp::client::GrpcChunkedBidiStream< 69 google::internal::federatedml::v2::ServerStreamMessage, 70 google::internal::federatedml::v2::ClientStreamMessage>* stream); 71 72 // Returns the client metadata from the most recent session call. 73 std::multimap<std::string, std::string> GetClientMetadata() const; 74 75 protected: 76 int32_t chunk_size_for_upload_; 77 int32_t max_pending_chunks_; 78 google::internal::federatedml::v2::CompressionLevel compression_level_; 79 absl::Notification session_done_; 80 81 private: 82 std::multimap<std::string, std::string> client_metadata_; 83 }; 84 85 } // namespace test 86 } // namespace client 87 } // namespace fcp 88 89 #endif // FCP_CLIENT_FAKE_SERVER_H_ 90