1 files changed, 41 insertions, 24 deletions

diff --git a/test/cpp/util/proto_file_parser.h b/test/cpp/util/proto_file_parser.h
index b442d77db9..eda3991e72 100644
--- a/test/cpp/util/proto_file_parser.h
+++ b/test/cpp/util/proto_file_parser.h
@@ -36,11 +36,9 @@
 
 #include <memory>
 
-#include <google/protobuf/compiler/importer.h>
-#include <google/protobuf/dynamic_message.h>
 #include <grpc++/channel.h>
 
-#include "src/compiler/config.h"
+#include "test/cpp/util/config_grpc_cli.h"
 #include "test/cpp/util/proto_reflection_descriptor_database.h"
 
 namespace grpc {
@@ -50,44 +48,63 @@ class ErrorPrinter;
 // Find method and associated request/response types.
 class ProtoFileParser {
  public:
-  // The given proto file_name will be searched in a source tree rooted from
-  // proto_path. The method could be a partial string such as Service.Method or
-  // even just Method. It will log an error if there is ambiguity.
-  ProtoFileParser(const grpc::string& proto_path, const grpc::string& file_name,
-                  const grpc::string& method);
-
+  // The parser will search proto files using the server reflection service
+  // provided on the given channel. The given protofiles in a source tree rooted
+  // from proto_path will also be searched.
   ProtoFileParser(std::shared_ptr<grpc::Channel> channel,
-                  const grpc::string& method);
+                  const grpc::string& proto_path,
+                  const grpc::string& protofiles);
+
   ~ProtoFileParser();
 
-  grpc::string GetFullMethodName() const { return full_method_name_; }
+  // The input method name in the following four functions could be a partial
+  // string such as Service.Method or even just Method. It will log an error if
+  // there is ambiguity.
+  // Full method name is in the form of Service.Method, it's good to be used in
+  // descriptor database queries.
+  grpc::string GetFullMethodName(const grpc::string& method);
+
+  // Formated method name is in the form of /Service/Method, it's good to be
+  // used as the argument of Stub::Call()
+  grpc::string GetFormatedMethodName(const grpc::string& method);
+
+  grpc::string GetSerializedProtoFromMethod(
+      const grpc::string& method, const grpc::string& text_format_proto,
+      bool is_request);
+
+  grpc::string GetTextFormatFromMethod(const grpc::string& method,
+                                       const grpc::string& serialized_proto,
+                                       bool is_request);
 
-  grpc::string GetSerializedProto(const grpc::string& text_format_proto,
-                                  bool is_request);
+  grpc::string GetSerializedProtoFromMessageType(
+      const grpc::string& message_type_name,
+      const grpc::string& text_format_proto);
 
-  grpc::string GetTextFormat(const grpc::string& serialized_proto,
-                             bool is_request);
+  grpc::string GetTextFormatFromMessageType(
+      const grpc::string& message_type_name,
+      const grpc::string& serialized_proto);
 
   bool HasError() const { return has_error_; }
 
   void LogError(const grpc::string& error_msg);
 
  private:
-  void InitProtoFileParser(
-      const grpc::string& method,
-      const std::vector<const google::protobuf::ServiceDescriptor*> services);
+  grpc::string GetMessageTypeFromMethod(const grpc::string& method,
+                                        bool is_request);
 
   bool has_error_;
   grpc::string request_text_;
-  grpc::string full_method_name_;
-  google::protobuf::compiler::DiskSourceTree source_tree_;
+  protobuf::compiler::DiskSourceTree source_tree_;
   std::unique_ptr<ErrorPrinter> error_printer_;
-  std::unique_ptr<google::protobuf::compiler::Importer> importer_;
-  std::unique_ptr<grpc::ProtoReflectionDescriptorDatabase> desc_db_;
-  std::unique_ptr<google::protobuf::DescriptorPool> desc_pool_;
-  std::unique_ptr<google::protobuf::DynamicMessageFactory> dynamic_factory_;
+  std::unique_ptr<protobuf::compiler::Importer> importer_;
+  std::unique_ptr<grpc::ProtoReflectionDescriptorDatabase> reflection_db_;
+  std::unique_ptr<protobuf::DescriptorPoolDatabase> file_db_;
+  std::unique_ptr<protobuf::DescriptorDatabase> desc_db_;
+  std::unique_ptr<protobuf::DescriptorPool> desc_pool_;
+  std::unique_ptr<protobuf::DynamicMessageFactory> dynamic_factory_;
   std::unique_ptr<grpc::protobuf::Message> request_prototype_;
   std::unique_ptr<grpc::protobuf::Message> response_prototype_;
+  std::vector<const protobuf::ServiceDescriptor*> service_desc_list_;
 };
 
 }  // namespace testing