Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.
24 changes: 24 additions & 0 deletions tests/deepsparse/pipelines/test_basic_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
Simple example and test of a dummy pipeline
"""

import asyncio
from typing import Dict

from pydantic import BaseModel
Expand All @@ -24,6 +25,7 @@
from deepsparse.operators import Operator
from deepsparse.routers import LinearRouter
from deepsparse.schedulers import OperatorScheduler
from deepsparse.utils.state import InferenceState


class IntSchema(BaseModel):
Expand Down Expand Up @@ -58,3 +60,25 @@ def test_run_simple_pipeline():
pipeline_output = AddThreePipeline(pipeline_input)

assert pipeline_output.value == 8


def test_run_async_simple_pipeline():
test_actually_ran = False

async def _actual_test():
nonlocal test_actually_ran

inference_state = InferenceState()
inference_state.create_state({})
pipeline_input = IntSchema(value=5)

pipeline_output = await AddThreePipeline.run_async(
pipeline_input, inference_state=inference_state
)

assert pipeline_output.value == 8

test_actually_ran = True

asyncio.run(_actual_test())
assert test_actually_ran