54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import asyncio
|
|
|
|
import cognee
|
|
from cognee.modules.pipelines.operations.run_tasks_base import run_tasks_base
|
|
from cognee.modules.pipelines.tasks.task import Task
|
|
from cognee.modules.users.methods import get_default_user
|
|
from cognee.infrastructure.databases.relational import create_db_and_tables
|
|
|
|
|
|
async def run_and_check_tasks():
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
def number_generator(num):
|
|
for i in range(num):
|
|
yield i + 1
|
|
|
|
async def add_one(nums):
|
|
for num in nums:
|
|
yield num + 1
|
|
|
|
async def multiply_by_two(nums):
|
|
yield nums[0] * 2
|
|
|
|
async def add_one_single(nums):
|
|
yield nums[0] + 1
|
|
|
|
await create_db_and_tables()
|
|
user = await get_default_user()
|
|
|
|
pipeline = run_tasks_base(
|
|
[
|
|
Task(number_generator),
|
|
Task(add_one, task_config={"batch_size": 5}),
|
|
Task(multiply_by_two, task_config={"batch_size": 1}),
|
|
Task(add_one_single),
|
|
],
|
|
data=10,
|
|
user=user,
|
|
)
|
|
|
|
results = [5, 7, 9, 11, 13, 15, 17, 19, 21, 23]
|
|
index = 0
|
|
async for result in pipeline:
|
|
assert result[0] == results[index], f"at {index = }: {result = } != {results[index] = }"
|
|
index += 1
|
|
|
|
|
|
def test_run_tasks():
|
|
asyncio.run(run_and_check_tasks())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_run_tasks()
|