diff --git a/tests/test_mlserver.py b/tests/test_mlserver.py
index 624134ba..934c9bb6 100644
--- a/tests/test_mlserver.py
+++ b/tests/test_mlserver.py
@@ -4,7 +4,7 @@
 import pytest
 from mlserver.settings import ModelSettings
 from mlserver.types import InferenceRequest, RequestInput
-from mlserver.utils import to_ndarray
+from mlserver.codecs import NumpyCodec
 from pytest_cases import fixture, parametrize_with_cases
 from pytest_cases.common_pytest_lazy_values import is_lazy
 
@@ -16,7 +16,11 @@
 
 @pytest.fixture
 def inference_request() -> InferenceRequest:
-    return InferenceRequest(inputs=[RequestInput(name="payload", shape=[4], data=[1, 2, 3, 4], datatype="FP32")])
+    return InferenceRequest(
+        inputs=[
+            RequestInput(name="payload", shape=[4], data=[1, 2, 3, 4], datatype="FP32")
+        ]
+    )
 
 
 @fixture
@@ -40,7 +44,9 @@ async def test_load(mlserver_runtime: InferenceRuntime):
     assert isinstance(mlserver_runtime._model, BaseModel)
 
 
-async def test_predict(mlserver_runtime: InferenceRuntime, inference_request: InferenceRequest):
+async def test_predict(
+    mlserver_runtime: InferenceRuntime, inference_request: InferenceRequest
+):
     # NOTE: pytest-cases doesn't wait for async fixtures
     # TODO: Raise issue in pytest-cases repo
     mlserver_runtime = await mlserver_runtime
@@ -48,7 +54,8 @@ async def test_predict(mlserver_runtime: InferenceRuntime, inference_request: In
 
     assert len(res.outputs) == 1
 
-    pipeline_input = to_ndarray(inference_request.inputs[0])
+    codec = NumpyCodec()
+    pipeline_input = codec.decode(inference_request.inputs[0])
     custom_model = copy.copy(mlserver_runtime._model)
     # Ensure direct call to class does not try to do remote
     custom_model.set_remote(False)
@@ -58,11 +65,14 @@ async def test_predict(mlserver_runtime: InferenceRuntime, inference_request: In
 
     pipeline_output = res.outputs[0].data
 
-    assert expected_output.tolist() == pipeline_output
+    assert expected_output.tolist() == pipeline_output.__root__
 
 
-async def test_load_wrapped_class(inference_pipeline_class, inference_request: InferenceRequest):
-    pipeline_input = to_ndarray(inference_request.inputs[0])
+async def test_load_wrapped_class(
+    inference_pipeline_class, inference_request: InferenceRequest
+):
+    codec = NumpyCodec()
+    pipeline_input = codec.decode(inference_request.inputs[0])
 
     inference_pipeline_class(pipeline_input)
     assert inference_pipeline_class.counter == 1