diff --git a/platypush/entities/__init__.py b/platypush/entities/__init__.py index 8ea774d0..6317419d 100644 --- a/platypush/entities/__init__.py +++ b/platypush/entities/__init__.py @@ -1,4 +1,6 @@ +from datetime import datetime, timedelta import logging +from threading import Event from typing import Collection, Optional from ._base import ( @@ -36,6 +38,23 @@ def init_entities_engine() -> EntitiesEngine: return _engine +def get_entities_engine(timeout: Optional[float] = None) -> EntitiesEngine: + """ + Return the running entities engine. + + :param timeout: Timeout in seconds (default: None). + """ + time_start = datetime.utcnow() + while not timeout or (datetime.utcnow() - time_start < timedelta(seconds=timeout)): + if _engine: + break + + Event().wait(1) + + assert _engine, 'The entities engine has not been initialized' + return _engine + + def publish_entities( entities: Collection[Entity], callback: Optional[EntitySavedCallback] = None ) -> None: