diff --git a/src/scriptworker/context.py b/src/scriptworker/context.py index 8a5d04df..5249bdd4 100644 --- a/src/scriptworker/context.py +++ b/src/scriptworker/context.py @@ -25,7 +25,7 @@ from taskcluster.aio import Queue from scriptworker import task_process -from scriptworker.exceptions import CoTError +from scriptworker.exceptions import CoTError, ScriptWorkerException from scriptworker.utils import load_json_or_yaml_from_url, makedirs, scriptworker_session log = logging.getLogger(__name__) @@ -235,13 +235,21 @@ def projects(self, projects: Optional[Dict[str, Any]]) -> None: @property def event_loop(self) -> asyncio.AbstractEventLoop: - """asyncio.BaseEventLoop: the running event loop. + """asyncio.BaseEventLoop: the event loop set on the context, or the running one. This fixture mainly exists to allow for overrides during unit tests. + Raises: + ScriptWorkerException: if no loop has been set and none is running. + """ if not self._event_loop: - self._event_loop = asyncio.get_event_loop() + try: + self._event_loop = asyncio.get_running_loop() + except RuntimeError as exc: + raise ScriptWorkerException( + "No event loop is running and none was set on the context. Set ``context.event_loop`` before use, or read it from a coroutine." + ) from exc return self._event_loop @event_loop.setter diff --git a/src/scriptworker/worker.py b/src/scriptworker/worker.py index 4f370b71..3283e3cc 100644 --- a/src/scriptworker/worker.py +++ b/src/scriptworker/worker.py @@ -236,14 +236,16 @@ def main(event_loop=None): Args: event_loop (asyncio.BaseEventLoop, optional): the event loop to use. - If None, use ``asyncio.get_event_loop()``. Defaults to None. + If None, create one with ``asyncio.new_event_loop()``. Defaults to None. """ context, credentials = get_context_from_cmdln(sys.argv[1:]) log.info("Scriptworker starting up at {} UTC".format(arrow.utcnow().format())) log.info("Worker FQDN: {}".format(socket.getfqdn())) cleanup(context) - context.event_loop = event_loop or asyncio.get_event_loop() + if event_loop is None: + event_loop = asyncio.new_event_loop() + context.event_loop = event_loop done = False diff --git a/tests/__init__.py b/tests/__init__.py index c6153728..cc8cec84 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -6,6 +6,7 @@ import json import os import sys +from contextlib import contextmanager import aiohttp import arrow @@ -46,6 +47,26 @@ def touch(path): print(path, file=fh, end="") +@contextmanager +def no_ambient_event_loop(): + """Run a block with no current event loop set for this thread. + + ``pytest-asyncio``'s auto mode installs a current event loop in the main + thread, so tests exercising the no-current-loop path have to remove it + first. The previous loop is restored on the way out, since sibling tests + run in a random order and expect to find it. + """ + try: + previous_loop = asyncio.get_event_loop() + except RuntimeError: + previous_loop = None + asyncio.set_event_loop(None) + try: + yield + finally: + asyncio.set_event_loop(previous_loop) + + class FakeResponse(aiohttp.client_reqrep.ClientResponse): """Integration tests allow us to test everything's hooked up to aiohttp correctly. When we don't want to actually hit an external url, have diff --git a/tests/conftest.py b/tests/conftest.py index d14b907d..4e13b45a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -171,5 +171,6 @@ def _craft_rw_context(tmp, cot_product, session, private=False): for rule in context.config["trusted_vcs_rules"]: rule["require_secret"] = True context.config["verbose"] = VERBOSE - context.event_loop = asyncio.new_event_loop() + # only reached from the async ``*_rw_context`` fixtures, so a loop is always running + context.event_loop = asyncio.get_running_loop() return context diff --git a/tests/test_context.py b/tests/test_context.py index eaad1f4f..db4f50eb 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -11,7 +11,7 @@ import pytest import scriptworker.context as swcontext -from scriptworker.exceptions import CoTError +from scriptworker.exceptions import CoTError, ScriptWorkerException # constants helpers and fixtures {{{1 @@ -155,12 +155,18 @@ def test_get_credentials(rw_context): assert rw_context.credentials == expected -def test_new_event_loop(mocker): - """The default rw_context.event_loop is from `asyncio.get_event_loop`""" - fake_loop = mock.MagicMock() - mocker.patch.object(asyncio, "get_event_loop", return_value=fake_loop) +def test_event_loop_no_running_loop(): + """`context.event_loop` raises when nothing is set and no loop is running.""" rw_context = swcontext.Context() - assert rw_context.event_loop is fake_loop + with pytest.raises(ScriptWorkerException): + rw_context.event_loop + + +@pytest.mark.asyncio +async def test_event_loop_prefers_running_loop(): + """`context.event_loop` returns the running loop when nothing is set.""" + rw_context = swcontext.Context() + assert rw_context.event_loop is asyncio.get_running_loop() def test_set_event_loop(mocker): diff --git a/tests/test_worker.py b/tests/test_worker.py index 3fefbdfb..410fa348 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -20,7 +20,18 @@ from scriptworker.exceptions import ScriptWorkerException, WorkerShutdownDuringTask from scriptworker.worker import RunTasks, do_run_task -from . import AT_LEAST_PY38, KILLED_SCRIPT, TIMEOUT_SCRIPT, create_async, create_finished_future, create_slow_async, create_sync, noop_async, noop_sync +from . import ( + AT_LEAST_PY38, + KILLED_SCRIPT, + TIMEOUT_SCRIPT, + create_async, + create_finished_future, + create_slow_async, + create_sync, + no_ambient_event_loop, + noop_async, + noop_sync, +) # constants helpers and fixtures {{{1 @@ -62,6 +73,35 @@ async def foo(arg, credentials): os.remove(tmp) +def test_main_no_ambient_event_loop(mocker, context): + """``main()`` runs with a loop it creates itself when none is current.""" + config = dict(context.config) + config["poll_interval"] = 1 + config["credentials"] = {"fake_creds": True} + + loops = [] + + async def foo(arg, credentials): + loops.append(arg.event_loop) + raise ScriptWorkerException("foo") + + _, tmp = tempfile.mkstemp() + try: + with open(tmp, "w") as fh: + json.dump(config, fh) + del config["credentials"] + mocker.patch.object(worker, "async_main", new=foo) + mocker.patch.object(sys, "argv", new=["x", tmp]) + with no_ambient_event_loop(): + with pytest.raises(ScriptWorkerException): + worker.main() + finally: + os.remove(tmp) + + assert len(loops) == 1 + assert isinstance(loops[0], asyncio.AbstractEventLoop) + + @pytest.mark.parametrize("running", (True, False)) def test_main_running_sigterm(mocker, context, running): """Test that sending SIGTERM causes the main loop to stop after the next