Skip to content

Commit

Permalink
fix: e2e tests
Browse files Browse the repository at this point in the history
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
  • Loading branch information
aarnphm committed Dec 6, 2022
1 parent be25b4f commit ceeb0ef
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions src/bentoml/testing/grpc/__init__.py
Expand Up @@ -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]
Expand All @@ -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.
Expand Down

0 comments on commit ceeb0ef

Please sign in to comment.