def test_retrying_rpc_exception(): first_call = [True] def predict_mock(request, timeout): inputs = tf.contrib.util.make_ndarray(request.inputs['images']) assert all([x == y for x, y in zip(inputs[0], [1, 2])]) if (first_call[0]): first_call[0] = False raise grpc.RpcError() return_data = np.asarray([[11, 22]]) return_tensor = tf.contrib.util.make_tensor_proto( return_data, types_pb2.DT_FLOAT, return_data.shape) result = mock.MagicMock() result.outputs = {"output_alias": return_tensor} return result stub_mock = mock.Mock() stub_mock.Predict = mock.MagicMock(side_effect=predict_mock) client = PredictionClient("localhost", 50051) client._get_grpc_stub = lambda: stub_mock result = client.score_numpy_array(np.asarray([[1, 2]], dtype='f')) assert all([x == y for x, y in zip(result[0], [11, 22])])
def test_score_numpy_array(): def predict_mock(request, timeout): inputs = tf.contrib.util.make_ndarray(request.inputs['images']) assert all([x == y for x, y in zip(inputs[0], [1, 2, 3])]) assert all([x == y for x, y in zip(inputs[1], [4, 5, 6])]) return_data = np.asarray([[11, 22, 33], [44, 55, 66]]) return_tensor = tf.contrib.util.make_tensor_proto( return_data, types_pb2.DT_FLOAT, return_data.shape) result = mock.MagicMock() result.outputs = {"output_alias": return_tensor} return result stub_mock = mock.Mock() stub_mock.Predict = mock.MagicMock(side_effect=predict_mock) client = PredictionClient("localhost", 50051) client._get_grpc_stub = lambda: stub_mock result = client.score_numpy_array( np.asarray([[1, 2, 3], [4, 5, 6]], dtype='f')) assert all([x == y for x, y in zip(result[0], [11, 22, 33])]) assert all([x == y for x, y in zip(result[1], [44, 55, 66])])
def test_score_image(): def predict_mock(request, timeout): inputs = request.inputs['images'].string_val assert inputs[0].decode('utf-8') == "abc" return_data = np.asarray([[1, 2, 3]]) return_tensor = tf.contrib.util.make_tensor_proto( return_data, types_pb2.DT_FLOAT, return_data.shape) result = mock.MagicMock() result.outputs = {"output_alias": return_tensor} return result stub_mock = mock.Mock() stub_mock.Predict = mock.MagicMock(side_effect=predict_mock) image_file_path = os.path.join(tempfile.mkdtemp(), "img.png") image_file = open(image_file_path, "w") image_file.write("abc") image_file.close() client = PredictionClient("localhost", 50051) client._get_grpc_stub = lambda: stub_mock result = client.score_image(image_file_path) assert all([x == y for x, y in zip(result, [1, 2, 3])])