baseimage-tutorial-framework/tfw/internals/ref_counter/test_ref_counter.py

# pylint: disable=redefined-outer-name
from dataclasses import dataclass
from textwrap import dedent
from os import mkfifo
from os.path import join
from signal import SIGINT
from subprocess import DEVNULL, Popen, PIPE
from tempfile import TemporaryDirectory

import pytest

from .ref_counter import RefCounter


@dataclass
class CounterContext:
    lockpath: str
    pipepath: str

    @property
    def source(self):
        return dedent(f'''\
            from time import sleep
            from atexit import register
            from ref_counter import RefCounter

            counter = RefCounter('{self.lockpath}')
            register(counter.teardown_instance)
            print(flush=True)
            while True:
                sleep(1)
        ''')


@pytest.fixture
def context():
    with TemporaryDirectory() as workdir:
        pipepath = join(workdir, 'test.pipe')
        mkfifo(pipepath)
        yield CounterContext(join(workdir, 'test.lock'), pipepath)


def test_increment_decrement(context):
    counter, processes = 0, []
    for _ in range(5):
        new_proc = Popen(['python3', '-c', context.source], stdout=PIPE, stderr=DEVNULL)
        new_proc.stdout.readline()
        processes.append(new_proc)
        counter += 1
    for proc in processes:
        with open(context.lockpath, 'r') as lock:
            assert lock.read() == str(counter)
            counter -= 1
        proc.send_signal(SIGINT)
        proc.wait()


def test_deallocate(context):
    state = False
    def trigger():
        nonlocal state
        state = True
    refcounters = []
    for _ in range(32):
        new_refc = RefCounter(context.lockpath)
        new_refc.deallocate = trigger
        refcounters.append(new_refc)
    for refc in refcounters:
        assert not state
        refc.teardown_instance()
    assert state