diff --git a/coco/client.py b/coco/client.py index cb5e62b..e0cd3bb 100644 --- a/coco/client.py +++ b/coco/client.py @@ -46,7 +46,7 @@ import yaml from aiohttp import ClientSession, ContentTypeError -from . import metric, result +from . import result from .config import DEFAULT_PORT, load_config from .endpoint import VALUE_TYPE @@ -325,7 +325,6 @@ def client_send_request( # Extract useful information from the click context host = ctx.obj["host"] port = ctx.obj["port"] - metrics_port = ctx.obj["coco_config"]["metrics_port"] url = f"http://{host}:{port}/{path}" @@ -346,11 +345,17 @@ def client_send_request( async def print_queue_size(metric_request_count): try: - q_size = await metric.get("coco_queue_length_total", metrics_port, host) - except RuntimeError as err: - if not isinstance(err, asyncio.CancelledError): - print(f"Couldn't get queue fill level from cocod: {err}") + async with ClientSession() as session: + async with session.get(f"http://{host}:{port}/qlen") as resp: + result = await resp.text() + q_size = int(result) + except asyncio.CancelledError: + # If cancelled, just return return + except (RuntimeError, ContentTypeError, KeyError) as err: + print(f"Couldn't get queue fill level from cocod: {err}") + return + print( f"\rThere are {int(q_size)} requests in the queue.", sep=" ", diff --git a/coco/core.py b/coco/core.py index 84e0a10..4e93320 100644 --- a/coco/core.py +++ b/coco/core.py @@ -117,7 +117,6 @@ def __init__( "reset-state", "save-state", "load-state", - "get-coco-config", "wait", } all_endpoints |= all_local_endpoints @@ -297,7 +296,10 @@ async def stop_slack_log(*_): self.sanic_app.register_listener(start_slack_log, "before_server_start") self.sanic_app.register_listener(stop_slack_log, "after_server_stop") + # non-endpoint routes. These take precedence over an identically named + # endpoint self.sanic_app.add_route(self._get_config, "/config", methods=["GET"]) + self.sanic_app.add_route(self._return_qlen, "/qlen", methods=["GET"]) self.sanic_app.add_route( self.external_endpoint, "/", methods=["GET", "POST"] @@ -483,6 +485,11 @@ async def _get_config(self, _): """ return response.json(self.config) + async def _return_qlen(self, _): + from sanic import text + + return text(str(self.redis_sync.llen("queue"))) + def _check_endpoint_links(self): def check(e): if e: diff --git a/tests/coco_runner.py b/tests/coco_runner.py index b1d7a7d..4d2f0de 100644 --- a/tests/coco_runner.py +++ b/tests/coco_runner.py @@ -35,6 +35,7 @@ def testN(coco_runner): import aiohttp import fakeredis import pytest +import redis import yaml from coco import config @@ -327,6 +328,22 @@ def start_daemon(self, *args): # even after it successfully fetches the port. sleep(0.2) + def redis_conn(self): + """Return a connection to the redis server. + + Each returns a new 'redis.Redis' connection. If the fakeredis server + isn't already running, it is started. + """ + + # Can't be called after stop() + if self._daemon_result: + raise RuntimeError("called after daemon stop.") + + # Start the redis server, if necessary + redis_port = self._start_redis() + + return redis.Redis(port=redis_port) + def call_endpoint( self, endpoint: str, @@ -363,6 +380,7 @@ def call_endpoint( # Can't be called after stop() if self._daemon_result: raise RuntimeError("called after daemon stop.") + self.start_daemon() # Build the URL diff --git a/tests/test_qlen.py b/tests/test_qlen.py new file mode 100644 index 0000000..372e650 --- /dev/null +++ b/tests/test_qlen.py @@ -0,0 +1,61 @@ +"""Test monitoring the coco queue length.""" + +import time + + +def waiting_callback(path, data): + """Reply only after waiting.""" + + # Wait... + time.sleep(1) + return data + + +def test_qlen(coco_runner): + """Test the /qlen route.""" + + # The "clean" targets don't wait. + # The "waiting" targets do. + coco_runner.add_targets("clean", 1) + coco_runner.add_targets("waiting", 1, callback=waiting_callback) + + coco_runner.add_endpoint("waiting", {"group": "waiting"}) + coco_runner.add_endpoint("clean", {"group": "clean"}) + coco_runner.start_daemon() + + # Check for clean daemon start + assert coco_runner.port is not None + + # Queue starts empty + response, text = coco_runner.call_endpoint("qlen") + assert response.status == 200 + assert int(text) == 0 + + # Connect to fakeredis + redis = coco_runner.redis_conn() + + # Directly add some stuff to the queue + NUM_REQ = 7 + for _ in range(NUM_REQ): + now = time.perf_counter() + tag = f"0-{now}" + redis.hset( + tag, + mapping={ + "method": "GET", + "endpoint": "waiting", + "request": "", + "params": "", + "received": now, + }, + ) + redis.rpush("queue", tag) + + # Check queue + response, text = coco_runner.call_endpoint("qlen") + assert response.status == 200 + assert int(text) == NUM_REQ - 1 # The qworker popped one + + # Can the client see this number? + result = coco_runner.client("clean") + assert "requests in the queue" in result.stdout