Internal change

PiperOrigin-RevId: 560406006
Intel-tensorflow · Aug 27, 2023 · ba44e06 · ba44e06
1 parent 07f9ba5
commit ba44e06
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/tensorflow_serving/model_servers/test_util/BUILD b/tensorflow_serving/model_servers/test_util/BUILD
@@ -31,6 +31,7 @@ cc_library(
         "//visibility:public",
     ],
     deps = [
+        "//tensorflow_serving/apis:logging_cc_proto",
         "//tensorflow_serving/apis:model_cc_proto",
         "//tensorflow_serving/config:model_server_config_cc_proto",
         "//tensorflow_serving/config:platform_config_cc_proto",

diff --git a/tensorflow_serving/model_servers/test_util/mock_server_core.h b/tensorflow_serving/model_servers/test_util/mock_server_core.h
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include <memory>
 #include <string>
+#include <utility>
 
 #include "base/logging.h"
 #include "google/protobuf/any.pb.h"
@@ -56,7 +57,9 @@ class MockServerCore : public ServerCore {
     return platform_config_map;
   }
 
-  static Options GetOptions(const PlatformConfigMap& platform_config_map) {
+  static Options GetOptions(
+      const PlatformConfigMap& platform_config_map,
+      std::unique_ptr<ServerRequestLogger> server_request_logger) {
     Options options;
     options.platform_config_map = platform_config_map;
     options.servable_state_monitor_creator =
@@ -71,13 +74,21 @@ class MockServerCore : public ServerCore {
            UniquePtrWithDeps<AspiredVersionsManager>* manager) -> Status {
       return Status();
     };
-    TF_CHECK_OK(
-        ServerRequestLogger::Create(nullptr, &options.server_request_logger));
+    if (server_request_logger != nullptr) {
+      options.server_request_logger = std::move(server_request_logger);
+    } else {
+      TF_CHECK_OK(
+          ServerRequestLogger::Create(nullptr, &options.server_request_logger));
+    }
     return options;
   }
 
   explicit MockServerCore(const PlatformConfigMap& platform_config_map)
-      : ServerCore(GetOptions(platform_config_map)) {}
+      : MockServerCore(platform_config_map, nullptr) {}
+  MockServerCore(const PlatformConfigMap& platform_config_map,
+                 std::unique_ptr<ServerRequestLogger> server_request_logger)
+      : ServerCore(GetOptions(platform_config_map,
+                              std::move(server_request_logger))) {}
 
   MOCK_METHOD(ServableStateMonitor*, servable_state_monitor, (),
               (const, override));