Skip to content

Commit

Permalink
Allow pinning a hook as the last of its stage
Browse files Browse the repository at this point in the history
  • Loading branch information
vkarak committed Sep 10, 2023
1 parent b117f2e commit 225d6dd
Show file tree
Hide file tree
Showing 6 changed files with 99 additions and 20 deletions.
4 changes: 2 additions & 2 deletions docs/regression_test_api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,9 @@ The use of this module is required only when creating new tests programmatically

.. autodecorator:: reframe.core.builtins.require_deps

.. autodecorator:: reframe.core.builtins.run_after(stage)
.. autodecorator:: reframe.core.builtins.run_after(stage, *, always_last=False)

.. autodecorator:: reframe.core.builtins.run_before(stage)
.. autodecorator:: reframe.core.builtins.run_before(stage, *, always_last=False)

.. autodecorator:: reframe.core.builtins.sanity_function

Expand Down
21 changes: 16 additions & 5 deletions reframe/core/builtins.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def final(fn):

# Hook-related builtins

def run_before(stage):
def run_before(stage, *, always_last=False):
'''Attach the decorated function before a certain pipeline stage.
The function will run just before the specified pipeline stage and it
Expand All @@ -47,14 +47,25 @@ def run_before(stage):
:param stage: The pipeline stage where this function will be attached to.
See :ref:`pipeline-hooks` for the list of valid stage values.
:param always_last: Run this hook always as the last one of the stage. In
a whole test hierarchy, only a single hook can be explicitly pinned at
the end of the same-stage sequence of hooks. If another hook is
declared as ``always_last`` in the same stage, an error will be
issued.
.. versionchanged:: 4.4
The ``always_last`` argument was added.
'''
return hooks.attach_to('pre_' + stage)

return hooks.attach_to('pre_' + stage, always_last)


def run_after(stage):
def run_after(stage, *, always_last=False):
'''Attach the decorated function after a certain pipeline stage.
This is analogous to :func:`~RegressionMixin.run_before`, except that the
This is analogous to :func:`run_before`, except that the
hook will execute right after the stage it was attached to. This decorator
also supports ``'init'`` as a valid ``stage`` argument, where in this
case, the hook will execute right after the test is initialized (i.e.
Expand All @@ -81,7 +92,7 @@ def __init__(self):
Add support for post-init hooks.
'''
return hooks.attach_to('post_' + stage)
return hooks.attach_to('post_' + stage, always_last)


require_deps = hooks.require_deps
Expand Down
9 changes: 5 additions & 4 deletions reframe/core/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,16 @@
import reframe.utility as util


def attach_to(phase):
def attach_to(phase, always_last):
'''Backend function to attach a hook to a given phase.
:meta private:
'''
def deco(func):
if hasattr(func, '_rfm_attach'):
func._rfm_attach.append(phase)
func._rfm_attach.append((phase, always_last))
else:
func._rfm_attach = [phase]
func._rfm_attach = [(phase, always_last)]

try:
# no need to resolve dependencies independently; this function is
Expand Down Expand Up @@ -124,6 +124,7 @@ def __init__(self, fn):
@property
def stages(self):
return self._rfm_attach
# return [stage for stage, _ in self._rfm_attach]

def __getattr__(self, attr):
return getattr(self.__fn, attr)
Expand Down Expand Up @@ -179,7 +180,7 @@ def add(self, v):
self.__hooks.discard(h)
self.__hooks.add(h)
elif hasattr(v, '_rfm_resolve_deps'):
v._rfm_attach = ['post_setup']
v._rfm_attach = [('post_setup', None)]
self.__hooks.add(Hook(v))

def update(self, other, *, denied_hooks=None):
Expand Down
30 changes: 25 additions & 5 deletions reframe/core/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,12 +179,32 @@ def disable_hook(self, hook_name):
@classmethod
def pipeline_hooks(cls):
ret = {}
last = {}
for hook in cls._rfm_hook_registry:
for stage in hook.stages:
try:
ret[stage].append(hook.fn)
except KeyError:
ret[stage] = [hook.fn]
for stage, always_last in hook.stages:
if always_last:
if stage in last:
hook_name = hook.__qualname__
pinned_name = last[stage].__qualname__
raise ReframeSyntaxError(
f'cannot pin hook {hook_name!r} as last '
f'of stage {stage!r} as {pinned_name!r} '
f'is already pinned last'
)

last[stage] = hook
else:
try:
ret[stage].append(hook.fn)
except KeyError:
ret[stage] = [hook.fn]

# Append the last hooks
for stage, hook in last.items():
try:
ret[stage].append(hook.fn)
except KeyError:
ret[stage] = [hook.fn]

return ret

Expand Down
8 changes: 4 additions & 4 deletions unittests/test_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ class Foo(MyMeta):
def hook_a(self):
pass

@run_before('compile')
@run_before('compile', always_last=True)
def hook_b(self):
pass

Expand All @@ -198,19 +198,19 @@ def hook_c(self):
pass

@classmethod
def hook_in_stage(cls, hook, stage):
def hook_in_stage(cls, hook, stage, always_last=False):
'''Assert that a hook is in a given registry stage.'''
for h in cls._rfm_hook_registry:
if h.__name__ == hook:
if stage in h.stages:
if (stage, always_last) in h.stages:
return True

break

return False

assert Foo.hook_in_stage('hook_a', 'post_setup')
assert Foo.hook_in_stage('hook_b', 'pre_compile')
assert Foo.hook_in_stage('hook_b', 'pre_compile', True)
assert Foo.hook_in_stage('hook_c', 'post_run')

class Bar(Foo):
Expand Down
47 changes: 47 additions & 0 deletions unittests/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -1162,6 +1162,53 @@ def foo(self):
assert test.pipeline_hooks() == {'post_setup': [MyTest.foo]}


def test_pinned_hooks():
@test_util.custom_prefix('unittests/resources/checks')
class X(rfm.RunOnlyRegressionTest):
@run_before('run', always_last=True)
def foo(self):
pass

class Y(X):
@run_before('run')
def bar(self):
pass

test = Y()
assert test.pipeline_hooks() == {'pre_run': [Y.bar, X.foo]}


def test_pinned_hooks_multiple_last():
@test_util.custom_prefix('unittests/resources/checks')
class X(rfm.RunOnlyRegressionTest):
@run_before('run', always_last=True)
def foo(self):
pass

class Y(X):
@run_before('run', always_last=True)
def bar(self):
pass

with pytest.raises(ReframeSyntaxError):
test = Y()


def test_pinned_hooks_multiple_last_inherited():
@test_util.custom_prefix('unittests/resources/checks')
class X(rfm.RunOnlyRegressionTest):
@run_before('run', always_last=True)
def foo(self):
pass

@run_before('run', always_last=True)
def bar(self):
pass

with pytest.raises(ReframeSyntaxError):
test = X()


def test_disabled_hooks(HelloTest, local_exec_ctx):
@test_util.custom_prefix('unittests/resources/checks')
class BaseTest(HelloTest):
Expand Down

0 comments on commit 225d6dd

Please sign in to comment.