Function test_default_batch_unbatch_stream

tests/unit/test_litapi.py:89–100

()

def test_default_batch_unbatch_stream():
    api = TestStreamAPIBatched(max_batch_size=4)
    api.request_timeout = 30
    api.stream = True
    api.pre_setup(spec=None)
    inputs = [1, 2, 3, 4]
    expected_output = [[0, 0, 0, 0], [1, 2, 3, 4], [2, 4, 6, 8], [3, 6, 9, 12]]
    output = api.batch(inputs)
    output = api.predict(output)
    for out in api.unbatch(output):
        expected = expected_output.pop(0)
        assert np.all(out == expected), f"Default unbatch should not change input {out} != {expected}"


def test_custom_batch_unbatch():