import asyncio

from neurovault.websocket_stream import EVENT_TYPES, EventStream


def test_publish_subscribe_event_stream():
    stream = EventStream()
    queue = stream.subscribe()

    loop = asyncio.new_event_loop()
    try:
        loop.run_until_complete(stream.publish("test.event", {"key": "value"}))
        event = loop.run_until_complete(asyncio.wait_for(queue.get(), timeout=1))
    finally:
        loop.close()

    assert event["type"] == "test.event"
    assert event["data"]["key"] == "value"


def test_event_type_catalog_contains_calendar_prediction():
    assert "prediction.calendar" in EVENT_TYPES
    assert "consolidation.entity_extracted" in EVENT_TYPES
