Testing Threaded Code

import unittest
from concurrent.futures import Executor, Future, as_completed


class InlineExecutor(Executor):
    """Runs tasks immediately in the caller's thread; returns completed futures."""

    def submit(self, fn, /, *args, **kwargs):
        fut: Future = Future()
        if fut.set_running_or_notify_cancel():
            try:
                fut.set_result(fn(*args, **kwargs))
            except BaseException as exc:
                fut.set_exception(exc)
        return fut


def collect(executor: Executor, fn, items):
    """Code under test: results and errors per item."""
    futures = {executor.submit(fn, x): x for x in items}
    ok, failed = {}, {}
    for fut in as_completed(futures):
        try:
            ok[futures[fut]] = fut.result()
        except Exception as exc:
            failed[futures[fut]] = type(exc).__name__
    return ok, failed


def parse(x: str) -> int:
    return int(x)


class CollectTests(unittest.TestCase):
    def test_errors_are_kept(self):
        ok, failed = collect(InlineExecutor(), parse, ["1", "x", "3"])
        self.assertEqual(ok, {"1": 1, "3": 3})
        self.assertEqual(failed, {"x": "ValueError"})

    def test_map_works_too(self):
        self.assertEqual(list(InlineExecutor().map(parse, ["4", "5"])), [4, 5])


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 2

import threading
import unittest
from concurrent.futures import ThreadPoolExecutor


class ConcurrencyTests(unittest.TestCase):
    def test_tasks_really_overlap(self):
        """Fails (BrokenBarrierError) if the pool runs tasks one at a time."""
        barrier = threading.Barrier(4, timeout=5)
        with ThreadPoolExecutor(max_workers=4) as ex:
            futures = [ex.submit(barrier.wait) for _ in range(4)]
            for f in futures:
                f.result(timeout=5)

    def test_second_writer_waits_for_lock(self):
        lock, inside, release = threading.Lock(), threading.Event(), threading.Event()
        order = []

        def first():
            with lock:
                inside.set()
                release.wait(5)
                order.append("first")

        def second():
            with lock:
                order.append("second")

        with ThreadPoolExecutor(max_workers=2) as ex:
            f1 = ex.submit(first)
            self.assertTrue(inside.wait(5))            # first holds the lock now
            f2 = ex.submit(second)
            self.assertFalse(f2.done())                # blocked on the lock
            release.set()
            f1.result(timeout=5)
            f2.result(timeout=5)
        self.assertEqual(order, ["first", "second"])


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 2

import threading
import unittest


class FakeClock:
    def __init__(self) -> None:
        self.now = 0.0
        self.sleeps = []
        self._lock = threading.Lock()

    def __call__(self) -> float:
        with self._lock:
            return self.now

    def sleep(self, seconds: float) -> None:
        with self._lock:
            self.sleeps.append(seconds)
            self.now += seconds


class TokenBucket:
    def __init__(self, rate_per_minute, burst, clock, sleep):
        self.rate, self.capacity, self.tokens = rate_per_minute / 60, float(burst), float(burst)
        self.clock, self.sleep = clock, sleep
        self.updated = clock()
        self.lock = threading.Lock()

    def acquire(self):
        while True:
            with self.lock:
                now = self.clock()
                self.tokens = min(self.capacity, self.tokens + (now - self.updated) * self.rate)
                self.updated = now
                if self.tokens >= 1:
                    self.tokens -= 1
                    return
                wait = (1 - self.tokens) / self.rate
            self.sleep(wait)


def retry(fn, attempts, base, sleep):
    for attempt in range(attempts):
        try:
            return fn()
        except ConnectionError:
            if attempt == attempts - 1:
                raise
            sleep(base * 2 ** attempt)                   # jitter omitted: inject an rng to test it


class TimeTests(unittest.TestCase):
    def test_bucket_allows_burst_then_rate(self):
        clock = FakeClock()
        bucket = TokenBucket(rate_per_minute=60, burst=5, clock=clock, sleep=clock.sleep)
        for _ in range(65):
            bucket.acquire()
        self.assertAlmostEqual(clock.now, 60.0)        # 5 immediate, then 1/s for 60 more

    def test_retry_backoff_schedule(self):
        clock = FakeClock()
        calls = iter([ConnectionError(), ConnectionError(), "ok"])

        def flaky():
            item = next(calls)
            if isinstance(item, Exception):
                raise item
            return item

        self.assertEqual(retry(flaky, attempts=5, base=0.5, sleep=clock.sleep), "ok")
        self.assertEqual(clock.sleeps, [0.5, 1.0])


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 2

import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from unittest import mock


class RateLimited(Exception):
    pass


def score_all(client, prompts, workers=8):
    """Code under test: one retry on RateLimited."""
    def one(p):
        try:
            return client.complete(p)
        except RateLimited:
            return client.complete(p)
    with ThreadPoolExecutor(max_workers=workers) as ex:
        return dict(zip(prompts, ex.map(one, prompts)))


class ClientTests(unittest.TestCase):
    def test_each_prompt_retried_once_on_429(self):
        seen = {}
        lock = threading.Lock()

        def fake_complete(prompt):
            with lock:
                seen[prompt] = seen.get(prompt, 0) + 1
                first = seen[prompt] == 1
            if prompt.endswith("7") and first:
                raise RateLimited(prompt)
            return prompt.upper()

        client = mock.Mock()
        client.complete.side_effect = fake_complete
        prompts = [f"p{i}" for i in range(20)]
        out = score_all(client, prompts)
        self.assertEqual(out["p17"], "P17")
        self.assertEqual(client.complete.call_count, 22)   # 20 + retries for p7 and p17
        self.assertEqual(seen["p7"], 2)


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 1

import itertools
import threading
import time
import unittest
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait


def window_run(fn, items, ex, max_in_flight):
    source = iter(items)
    pending = {ex.submit(fn, x) for x in itertools.islice(source, max_in_flight)}
    results = []
    while pending:
        done, pending = wait(pending, return_when=FIRST_COMPLETED)
        results += [f.result() for f in done]
        pending |= {ex.submit(fn, x) for x in itertools.islice(source, len(done))}
    return results


class BoundTests(unittest.TestCase):
    def test_never_more_than_window_pending(self):
        lock, live, peak = threading.Lock(), [0], [0]
        pulled = []

        def lazy():
            for i in range(200):
                pulled.append(i)
                with lock:                               # items pulled but not yet finished
                    live[0] += 1
                    peak[0] = max(peak[0], live[0])
                yield i

        def work(i):
            time.sleep(0.001)
            with lock:
                live[0] -= 1
            return i

        with ThreadPoolExecutor(max_workers=4) as ex:
            out = window_run(work, lazy(), ex, max_in_flight=8)
        self.assertEqual(sorted(out), list(range(200)))
        self.assertLessEqual(peak[0], 8)

    def test_cancel_futures_drops_queued_work(self):
        gate, ran = threading.Event(), []

        def task(i):
            gate.wait(5)
            ran.append(i)

        ex = ThreadPoolExecutor(max_workers=2)
        futures = [ex.submit(task, i) for i in range(50)]
        gate.set()
        ex.shutdown(wait=True, cancel_futures=True)
        cancelled = sum(f.cancelled() for f in futures)
        self.assertEqual(cancelled + len(ran), 50)
        self.assertGreater(cancelled, 0)

    def test_ctrl_c_path_sets_stop_and_cancels(self):
        stop = threading.Event()

        def task(i):
            stop.wait(5)
            return i

        ex = ThreadPoolExecutor(max_workers=2)
        futures = [ex.submit(task, i) for i in range(20)]
        try:
            raise KeyboardInterrupt                      # what the main thread's wait() would raise
        except KeyboardInterrupt:
            stop.set()
            ex.shutdown(wait=True, cancel_futures=True)
        self.assertTrue(all(f.done() for f in futures))
        self.assertGreaterEqual(sum(f.cancelled() for f in futures), 18)


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 3

import contextvars
import unittest
from concurrent.futures import ThreadPoolExecutor

request_id = contextvars.ContextVar("request_id", default="-")


class ContextExecutor(ThreadPoolExecutor):
    def submit(self, fn, /, *args, **kwargs):
        return super().submit(contextvars.copy_context().run, fn, *args, **kwargs)


class ContextTests(unittest.TestCase):
    def test_request_id_reaches_workers_and_does_not_leak(self):
        with ContextExecutor(max_workers=1) as ex:
            token = request_id.set("req-1")
            try:
                self.assertEqual(ex.submit(request_id.get).result(timeout=5), "req-1")
                ex.submit(request_id.set, "inner").result(timeout=5)
                self.assertEqual(ex.submit(request_id.get).result(timeout=5), "req-1")
            finally:
                request_id.reset(token)
            self.assertEqual(list(ex.map(lambda _: request_id.get(), [1, 2])), ["-", "-"])


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 1

import threading
import time
import unittest
from concurrent.futures import ThreadPoolExecutor


class LeakCheckMixin:
    def setUp(self):
        self._threads_before = set(threading.enumerate())
        super().setUp()

    def tearDown(self):
        super().tearDown()
        deadline = time.monotonic() + 2
        while True:
            leaked = [t for t in threading.enumerate() if t not in self._threads_before and t.is_alive()]
            if not leaked or time.monotonic() > deadline:
                break
            time.sleep(0.01)
        if leaked:
            raise AssertionError(f"leaked threads: {sorted(t.name for t in leaked)}")


def leaky_job():
    ex = ThreadPoolExecutor(max_workers=2, thread_name_prefix="leaky")
    ex.submit(time.sleep, 0.01).result()
    return ex                                            # never shut down


class Jobs(LeakCheckMixin, unittest.TestCase):
    def test_clean_job(self):
        with ThreadPoolExecutor(max_workers=2, thread_name_prefix="clean") as ex:
            ex.submit(time.sleep, 0.01).result()


class LeakyJobs(LeakCheckMixin, unittest.TestCase):
    def test_leaky_job_is_caught(self):
        self.ex = leaky_job()

    def tearDown(self):
        with self.assertRaises(AssertionError):
            super().tearDown()
        self.ex.shutdown()


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 2

import threading
import time

import pytest


@pytest.fixture(autouse=True)
def no_leaked_threads():
    before = set(threading.enumerate())
    yield
    deadline = time.monotonic() + 2
    leaked = []
    while time.monotonic() < deadline:
        leaked = [t for t in threading.enumerate() if t not in before and t.is_alive()]
        if not leaked:
            return
        time.sleep(0.01)
    pytest.fail(f"leaked threads: {sorted(t.name for t in leaked)}")

import sys
import threading
import time
import unittest
from concurrent.futures import ThreadPoolExecutor


class Counter:
    def __init__(self, locked: bool) -> None:
        self.n = 0
        self.locked = locked
        self.lock = threading.Lock()

    def add(self) -> None:
        if self.locked:
            with self.lock:
                self.n += 1
        else:
            n = self.n                                  # explicit read-modify-write
            time.sleep(0)                               # stands in for real work between read and write
            self.n = n + 1


def hammer(counter: Counter, threads: int = 8, per_thread: int = 500) -> int:
    def run():
        for _ in range(per_thread):
            counter.add()
    with ThreadPoolExecutor(max_workers=threads) as ex:
        for f in [ex.submit(run) for _ in range(threads)]:
            f.result(timeout=30)
    return counter.n


class RaceTests(unittest.TestCase):
    def setUp(self):
        self.old = sys.getswitchinterval()
        sys.setswitchinterval(1e-6)

    def tearDown(self):
        sys.setswitchinterval(self.old)

    def test_locked_counter_is_exact(self):
        for _ in range(5):
            self.assertEqual(hammer(Counter(locked=True)), 4000)

    def test_unlocked_counter_loses_updates(self):
        lost = any(hammer(Counter(locked=False)) < 4000 for _ in range(5))
        self.assertTrue(lost)                            # demonstrates the race the linter's TPE006 flags


result = unittest.main(argv=["tests"], exit=False, verbosity=0).result
print(result.wasSuccessful(), result.testsRun)          # True 2