more coverage (#63)

* more coverage

* Add stack sampler to tests

Author: lsbardel

docs/reference/workers.md

@@ -1,10 +1,10 @@
 # Workers
 
-Workers are the main building block for asynchronous programming with `aio-fluid`. They are responsible for running tasks and managing their lifecycle.
-There are several worker classes which can be imported from `fluid.utils.worker`:
+Workers are the main building block for asynchronous programming with `aio-fluid`. They are responsible for running asynchronous tasks and managing their lifecycle.
+There are several worker classes which can be imported from `fluid.utils.worker`, and they aall derive from the abstract `fluid.utils.worker.Worker` class.
 
 ```python
-from fluid.utils.worker import StoppingWorker
+from fluid.utils.worker import Worker
 ```
 
 ::: fluid.utils.worker.WorkerState

fluid/utils/data.py

@@ -3,10 +3,6 @@
 DEFAULT_SKIP_VALUES = frozenset((None,))
 
 
-def reversed_dict(d: dict) -> dict:  # type: ignore [type-arg]
-    return {v: k for k, v in d.items()}
-
-
 def compact_dict(
     *args: Iterable[Any],
     skip_values: set[Any] | frozenset[Any] = DEFAULT_SKIP_VALUES,

fluid/utils/errors.py

@@ -26,3 +26,7 @@ def __str__(self) -> str:
 
 class WorkerStartError(FluidError):
     """Worker start error"""
+
+
+class FlamegraphError(FluidError):
+    pass

fluid/utils/executor.py

@@ -28,7 +28,7 @@ class Sampler:
     """
 
     interval: float = 0.005
-    _started = None
+    _started: float | None = None
     _stack_counts: defaultdict[str, int] = field(
         default_factory=lambda: defaultdict(int)
     )

fluid/utils/stacksampler.py

@@ -1,5 +1,4 @@
 import json
-import os
 import uuid
 from typing import Any
 
@@ -35,28 +34,14 @@ def as_uuid(uid: Any) -> str | None:
     return None
 
 
-def nice_env_str(space: int = 4, trim_length: int = 100) -> str:
-    lt = max(len(k) for k in os.environ) + space
-    values = []
-    for key, value in os.environ.items():
-        if len(value) > trim_length + 3:
-            value = f"{value[:trim_length]}..."
-        k = f"{key}:".ljust(lt)
-        values.append(f"{k}{value}")
-    return "\n".join(values)
-
-
 def nice_json(data: Any) -> str:
     if not isinstance(data, str):
         return json.dumps(data, indent=4)
     return data
 
 
 def trim_docstring(docstring: str) -> str:
-    """Uniformly trims leading/trailing whitespace from docstrings.
-    Based on
-    http://www.python.org/peps/pep-0257.html#handling-docstring-indentation
-    """
+    """Uniformly trims leading/trailing whitespace from doc-strings"""
     if not docstring or not docstring.strip():
         return ""
     # Convert tabs to spaces and split into lines

fluid/utils/text.py

@@ -70,14 +70,7 @@ async def shutdown(self) -> None:
     async def wait_for_shutdown(self) -> None:
         if self.task is None:
             return
-        try:
-            await self.task
-        except asyncio.CancelledError:
-            if self.force_shutdown:
-                # we are shutting down, this is expected
-                pass
-            else:
-                raise
+        await self.task
 
     async def gracefully_shutdown(self) -> None:
         self.started_shutdown = True
@@ -193,9 +186,6 @@ def is_stopping(self) -> bool:
     def is_stopped(self) -> bool:
         return self._worker_state in (WorkerState.STOPPED, WorkerState.FORCE_STOPPED)
 
-    def is_task(self) -> bool:
-        return self._worker_task_runner is not None
-
     def gracefully_stop(self) -> None:
         if self.is_running():
             self._worker_state = WorkerState.STOPPING
@@ -245,6 +235,11 @@ async def shutdown(self) -> None:
             await self._worker_task_runner.shutdown()
 
     async def wait_for_shutdown(self) -> None:
+        """Wait for the worker to stop
+
+        This method will wait for the worker to stop running, but doesn't
+        try to gracefully stop it nor force shutdown.
+        """
         if self._worker_task_runner is not None:
             await self._worker_task_runner.wait_for_shutdown()
 

fluid/utils/worker.py

@@ -1,5 +1,5 @@
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, cast
+from typing import AsyncIterator, Iterator, cast
 
 import pytest
 from fastapi import FastAPI
@@ -11,6 +11,7 @@
 from fluid.scheduler.broker import RedisTaskBroker
 from fluid.scheduler.endpoints import get_task_manger, setup_fastapi
 from fluid.tools_fastapi import backdoor
+from fluid.utils.stacksampler import Sampler
 from tests.scheduler.tasks import TaskClient
 
 
@@ -25,6 +26,14 @@ def redis_broker(task_manager: TaskManager) -> RedisTaskBroker:
     return cast(RedisTaskBroker, task_manager.broker)
 
 
+@pytest.fixture(scope="module", autouse=True)
+def sampler() -> Iterator[Sampler]:
+    sampler = Sampler()
+    sampler.start()
+    yield sampler
+    sampler.stop()
+
+
 @pytest.fixture(scope="module")
 async def task_app():
     task_manager = tasks.task_scheduler(max_concurrent_tasks=2, schedule_tasks=False)

tests/scheduler/conftest.py

@@ -12,6 +12,7 @@
     TaskState,
 )
 from fluid.scheduler.errors import UnknownTaskError
+from fluid.utils.stacksampler import Sampler
 from fluid.utils.waiter import wait_for
 
 pytestmark = pytest.mark.asyncio(loop_scope="module")
@@ -116,3 +117,9 @@ async def test_async_handler(task_scheduler: TaskScheduler) -> None:
     assert handler.task_run
     assert handler.task_run.state == TaskState.running
     assert task_scheduler.unregister_async_handler("running.test") is handler
+
+
+def test_sampler(sampler: Sampler) -> None:
+    assert sampler.started
+    stats = sampler.stats()
+    assert stats

tests/scheduler/test_scheduler.py

@@ -1,6 +1,9 @@
 import asyncio
 from dataclasses import dataclass, field
 
+import pytest
+
+from fluid.utils.errors import WorkerStartError
 from fluid.utils.worker import QueueConsumerWorker, Worker, Workers, WorkerState
 
 
@@ -10,6 +13,18 @@ async def run(self):
             await asyncio.sleep(0.1)
 
 
+class BadWorker2(Worker):
+    async def run(self):
+        await asyncio.sleep(0.1)
+        raise RuntimeError("I'm so bad")
+
+
+class BadWorker3(Worker):
+    async def run(self):
+        await asyncio.sleep(0.1)
+        raise asyncio.CancelledError
+
+
 @dataclass
 class Waiter:
     waiter: asyncio.Future = field(
@@ -35,10 +50,37 @@ async def test_consumer() -> None:
     assert runner.is_stopped()
 
 
+async def test_metadata() -> None:
+    worker = BadWorker(stopping_grace_period=2)
+    assert worker.num_workers == 1
+    assert list(worker.workers()) == [worker]
+    await worker.wait_for_shutdown()
+    assert worker.worker_state is WorkerState.INIT
+
+
 async def test_force_shutdown() -> None:
     worker = BadWorker(stopping_grace_period=2)
     await worker.startup()
+    with pytest.raises(WorkerStartError):
+        await worker.startup()
     assert worker.is_running()
     await worker.shutdown()
     assert worker.is_stopped()
     assert worker.worker_state is WorkerState.FORCE_STOPPED
+
+
+async def test_exeception2() -> None:
+    worker = BadWorker2()
+    await worker.startup()
+    with pytest.raises(RuntimeError) as exc:
+        await worker.wait_for_shutdown()
+    assert str(exc.value) == "I'm so bad"
+    assert worker.is_stopped()
+
+
+async def test_exeception3() -> None:
+    worker = BadWorker3()
+    await worker.startup()
+    with pytest.raises(asyncio.CancelledError):
+        await worker.wait_for_shutdown()
+    assert worker.is_stopped()