diff --git a/tests/deepsparse/pipelines/test_basic_pipeline.py b/tests/deepsparse/pipelines/test_basic_pipeline.py index c98ffd538c..4e7ef79399 100644 --- a/tests/deepsparse/pipelines/test_basic_pipeline.py +++ b/tests/deepsparse/pipelines/test_basic_pipeline.py @@ -16,6 +16,7 @@ Simple example and test of a dummy pipeline """ +import asyncio from typing import Dict from pydantic import BaseModel @@ -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): @@ -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