def test_load_model(self):
client = ModelsClient()
name = "simple-nn-model"
version = "1"
model = client.load_model(name, version)
x = np.float32([[1.0], [2.0]])
y = model.predict(x)
assert y.shape[0] == 2
assert y.shape[1] == 1
def test_load_model(self, mocker):
mock_method = mocker.patch.object(mlflow.pyfunc, "load_model")
mock_method.return_value = mlflow.pytorch._PyTorchWrapper(
LinearNNModel())
client = ModelsClient()
name = "simple-nn-model"
version = "1"
model = client.load_model(name, version)
mock_method.assert_called_once_with(
model_uri="models:/simple-nn-model/1")
x = np.float32([[1.0], [2.0]])
y = model.predict(x)
assert y.shape[0] == 2
assert y.shape[1] == 1