From ceeb0ef452ec344004a21d64d6e420577c91fbcd Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Mon, 5 Dec 2022 17:30:10 -0800
Subject: [PATCH] fix: e2e tests

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
---
 src/bentoml/testing/grpc/__init__.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/bentoml/testing/grpc/__init__.py b/src/bentoml/testing/grpc/__init__.py
index b28922f747b..7367460f12d 100644
--- a/src/bentoml/testing/grpc/__init__.py
+++ b/src/bentoml/testing/grpc/__init__.py
@@ -54,17 +54,24 @@ def create_bento_servicer(
         ) from None
 
 
-def randomize_pb_ndarray(shape: tuple[int, ...]) -> pb.NDArray:
+def randomize_pb_ndarray(
+    shape: tuple[int, ...], protocol_version: str = LATEST_PROTOCOL_VERSION
+) -> pb.NDArray:
+    pb, _ = import_generated_stubs(protocol_version)
     arr: NDArray[np.float32] = t.cast("NDArray[np.float32]", np.random.rand(*shape))
     return pb.NDArray(
         shape=list(shape), dtype=pb.NDArray.DTYPE_FLOAT, float_values=arr.ravel()
     )
 
 
-def make_pb_ndarray(arr: NDArray[t.Any]) -> pb.NDArray:
+def make_pb_ndarray(
+    arr: NDArray[t.Any], protocol_version: str = LATEST_PROTOCOL_VERSION
+) -> pb.NDArray:
     from bentoml._internal.io_descriptors.numpy import npdtype_to_dtypepb_map
     from bentoml._internal.io_descriptors.numpy import npdtype_to_fieldpb_map
 
+    pb, _ = import_generated_stubs(protocol_version)
+
     try:
         fieldpb = npdtype_to_fieldpb_map()[arr.dtype]
         dtypepb = npdtype_to_dtypepb_map()[arr.dtype]
@@ -91,7 +98,7 @@ async def async_client_call(
     assert_code: grpc.StatusCode | None = None,
     assert_details: str | None = None,
     assert_trailing_metadata: aio.Metadata | None = None,
-    protocol_version: str = "v1",
+    protocol_version: str = LATEST_PROTOCOL_VERSION,
 ) -> pb.Response | None:
     """
     Invoke a given API method via a client.