def test_serve_ppo(self, mock_shutdown, mock_init, mock_sleep):
mock_sleep.side_effect = KeyboardInterrupt()
with mock.patch("blaze.serve.server.Server",
new=MockServer()) as mock_server:
with tempfile.NamedTemporaryFile() as model_location:
serve([
"--model",
"PPO",
"--port",
"5678",
"--host",
"127.0.0.1",
"--max_workers",
"16",
model_location.name,
])
assert mock_server.args[0].host == "127.0.0.1"
assert mock_server.args[0].port == 5678
assert mock_server.args[0].max_workers == 16
assert mock_server.set_policy_service_args[
0].saved_model.cls == PPOAgent
assert mock_server.set_policy_service_args[
0].saved_model.location == model_location.name
assert mock_server.start_called
assert mock_server.stop_called
def test_serve_invalid_model_location(self):
with pytest.raises(IOError):
serve(["--model", "APEX", "/non/existent/file"])
def test_serve_exits_with_invalid_arguments(self):
with pytest.raises(SystemExit):
serve([])