49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
from asyncio import sleep
|
|
|
|
import pytest
|
|
from great_ai import GreatAI
|
|
|
|
|
|
def test_process_batch() -> None:
|
|
@GreatAI.create
|
|
def f(x: int) -> int:
|
|
return x + 2
|
|
|
|
assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]
|
|
|
|
|
|
def test_process_batch_unpacked() -> None:
|
|
@GreatAI.create
|
|
def f(a: int, b: str) -> str:
|
|
return b * a
|
|
|
|
assert [
|
|
v.output
|
|
for v in f.process_batch(
|
|
[(2, "aa"), (1, "fa"), (3, "b")], unpack_arguments=True
|
|
)
|
|
] == ["aaaa", "fa", "bbb"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_batch_async() -> None:
|
|
@GreatAI.create
|
|
async def f(x: int) -> int:
|
|
await sleep(0.2)
|
|
return x + 2
|
|
|
|
assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_batch_async_unpacked() -> None:
|
|
@GreatAI.create
|
|
async def f(a: int, b: str) -> str:
|
|
return b * a
|
|
|
|
assert [
|
|
v.output
|
|
for v in f.process_batch(
|
|
[(2, "aa"), (1, "fa"), (3, "b")], unpack_arguments=True
|
|
)
|
|
] == ["aaaa", "fa", "bbb"]
|