mirror of
https://github.com/Monadical-SAS/reflector.git
synced 2025-12-20 20:29:06 +00:00
server: add BroadcastProcessor tests
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
from typing import Any, Union
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
@@ -211,12 +211,16 @@ class BroadcastProcessor(Processor):
|
||||
This processor does not guarantee that the output is in order.
|
||||
|
||||
This processor connect all the output of the processors to the input of
|
||||
the next processor.
|
||||
the next processor; so the next processor must be able to accept different
|
||||
types of input.
|
||||
"""
|
||||
|
||||
def __init__(self, processors: Processor):
|
||||
super().__init__()
|
||||
self.processors = processors
|
||||
self.INPUT_TYPE = processors[0].INPUT_TYPE
|
||||
output_types = set([processor.OUTPUT_TYPE for processor in processors])
|
||||
self.OUTPUT_TYPE = Union[tuple(output_types)]
|
||||
|
||||
def set_pipeline(self, pipeline: "Pipeline"):
|
||||
super().set_pipeline(pipeline)
|
||||
|
||||
46
server/tests/test_processors_broadcast.py
Normal file
46
server/tests/test_processors_broadcast.py
Normal file
@@ -0,0 +1,46 @@
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processor_broadcast():
|
||||
from reflector.processors.base import Processor, BroadcastProcessor, Pipeline
|
||||
|
||||
class TestProcessor(Processor):
|
||||
INPUT_TYPE = str
|
||||
OUTPUT_TYPE = str
|
||||
|
||||
def __init__(self, name, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.name = name
|
||||
|
||||
async def _push(self, data):
|
||||
data = data + f":{self.name}"
|
||||
await self.emit(data)
|
||||
|
||||
processors = [
|
||||
TestProcessor("A"),
|
||||
BroadcastProcessor(
|
||||
processors=[
|
||||
TestProcessor("B"),
|
||||
TestProcessor("C"),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
events = []
|
||||
|
||||
async def on_event(event):
|
||||
events.append(event)
|
||||
|
||||
pipeline = Pipeline(*processors)
|
||||
pipeline.on(on_event)
|
||||
await pipeline.push("test")
|
||||
await pipeline.flush()
|
||||
|
||||
assert len(events) == 3
|
||||
assert events[0].processor == "A"
|
||||
assert events[0].data == "test:A"
|
||||
assert events[1].processor == "B"
|
||||
assert events[1].data == "test:A:B"
|
||||
assert events[2].processor == "C"
|
||||
assert events[2].data == "test:A:C"
|
||||
Reference in New Issue
Block a user