mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-09-28 09:35:42 +00:00
516db4cb52
* original-component * stub * fix implementation * fix tests * review feedback * review feedback * upgrade canals * upgrade canals * upgrade canals to fix pipeline test * remove requests_with_retry * feedback
57 lines
1.4 KiB
Python
57 lines
1.4 KiB
Python
from typing import Dict, Any
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
|
|
from haystack.preview import Pipeline, component, NoSuchStoreError, ComponentInput, ComponentOutput
|
|
|
|
|
|
class MockStore:
|
|
...
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_pipeline_store_api():
|
|
store_1 = MockStore()
|
|
store_2 = MockStore()
|
|
pipe = Pipeline()
|
|
|
|
pipe.add_store(name="first_store", store=store_1)
|
|
pipe.add_store(name="second_store", store=store_2)
|
|
|
|
assert pipe.list_stores() == ["first_store", "second_store"]
|
|
|
|
assert pipe.get_store("first_store") == store_1
|
|
assert pipe.get_store("second_store") == store_2
|
|
with pytest.raises(NoSuchStoreError):
|
|
pipe.get_store("third_store")
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_pipeline_stores_in_params():
|
|
store_1 = MockStore()
|
|
store_2 = MockStore()
|
|
|
|
@component
|
|
class MockComponent:
|
|
@dataclass
|
|
class Input(ComponentInput):
|
|
value: int
|
|
stores: Dict[str, Any]
|
|
|
|
@dataclass
|
|
class Output(ComponentOutput):
|
|
value: int
|
|
|
|
def run(self, data: Input) -> Output:
|
|
assert data.stores == {"first_store": store_1, "second_store": store_2}
|
|
return MockComponent.Output(value=data.value)
|
|
|
|
pipe = Pipeline()
|
|
pipe.add_component("component", MockComponent())
|
|
|
|
pipe.add_store(name="first_store", store=store_1)
|
|
pipe.add_store(name="second_store", store=store_2)
|
|
|
|
pipe.run(data={"component": MockComponent.Input(value=1)})
|