Simple Barrier demo

barrier_example.py

# barrier_example.py
"""Example usage of barrier."""

import multiprocessing as mp
import itertools

def worker(barrier:mp.Barrier, q:mp.Queue=None):
    print(mp.current_process().name, 'waiting')
    i = barrier.wait()
    if i == 0:
        print('pass')
    if q:
        q.put(i)
    else:
        return i


def main(n:int=4):
    """Collect some values from workers using a barrier."""
    barrier = mp.Barrier(n)
    q = mp.Queue()
    ps = [mp.Process(target=worker, args=(barrier,q)) for _ in range(n)]
    [p.start() for p in ps]
    for _ in range(n):
        yield q.get()

def main_map(n:int=4):
    """When using a pool, we need a manager to share the barrier."""
    with mp.Manager() as m:
        barrier = m.Barrier(n)
        with mp.Pool() as pool:
            results = pool.map(worker, [barrier] * 4)
            return results


# test_barrier_example.py
"""Tests for barrier_example.py."""
import pytest
import mock

def test_main():
    assert set(main()) == set(range(4))

def test_main_map():
    assert set(main_map()) == set(range(4))
    
pytest.main(['-s'])