import pytest import redis from typing import List import subprocess from subprocess import Popen, check_output, TimeoutExpired from os import environ, path, set_blocking from io import BufferedReader from huey.exceptions import HueyException from selfprivacy_api.utils.huey import huey, immediate, HUEY_DATABASE_NUMBER from selfprivacy_api.utils.redis_pool import RedisPool, REDIS_SOCKET @huey.task() def sum(a: int, b: int) -> int: return a + b def reset_huey_storage(): huey.storage = huey.create_storage() def flush_huey_redis_forcefully(): url = RedisPool.connection_url(HUEY_DATABASE_NUMBER) pool = redis.ConnectionPool.from_url(url, decode_responses=True) connection = redis.Redis(connection_pool=pool) connection.flushdb() # TODO: may be useful in other places too, move to utils/ tests common if using it somewhere def read_all_ready_output(stream: BufferedReader) -> str: set_blocking(stream.fileno(), False) output: List[bytes] = [] while True: line = stream.readline() raise ValueError(line) if line == b"": break else: output.append(line) set_blocking(stream.fileno(), True) result = b"".join(output) return result.decode("utf-8") @pytest.fixture() def not_immediate(): assert environ["TEST_MODE"] == "true" old_immediate = huey.immediate environ["HUEY_QUEUES_FOR_TESTS"] = "Yes" huey.immediate = False assert huey.immediate is False yield del environ["HUEY_QUEUES_FOR_TESTS"] huey.immediate = old_immediate assert huey.immediate == old_immediate @pytest.fixture() def huey_socket_consumer(not_immediate): """ Same as above, but with socketed redis """ flush_huey_redis_forcefully() command = ["huey_consumer.py", "selfprivacy_api.task_registry.huey"] # First assert that consumer does not fail by itself # Idk yet how to do it more elegantly try: check_output(command, timeout=2) except TimeoutExpired: pass # Then open it for real consumer_handle = Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) assert path.exists(REDIS_SOCKET) yield consumer_handle consumer_handle.kill() def test_huey_over_redis_socket(huey_socket_consumer): assert huey.immediate is False assert immediate() is False assert "unix" in RedisPool.connection_url(HUEY_DATABASE_NUMBER) try: assert ( RedisPool.connection_url(HUEY_DATABASE_NUMBER) in huey.storage_kwargs.values() ) except AssertionError: raise ValueError( "our test-side huey does not connect over socket: ", huey.storage_kwargs ) result = sum(2, 5) try: assert result(blocking=True, timeout=100) == 7 except HueyException as error: if "timed out" in str(error): output = read_all_ready_output(huey_socket_consumer.stdout) errorstream = read_all_ready_output(huey_socket_consumer.stderr) raise TimeoutError( f"Huey timed out: {str(error)}", f"Consumer output: {output}", f"Consumer errorstream: {errorstream}", ) else: raise error @pytest.mark.xfail(reason="cannot yet schedule with sockets for some reason") def test_huey_schedule(huey_queues_socket): # We do not schedule tasks anywhere, but concerning that it fails. sum.schedule((2, 5), delay=10) try: assert len(huey.scheduled()) == 1 except AssertionError: raise ValueError("have wrong amount of scheduled tasks", huey.scheduled())