diff --git a/belay/device.py b/belay/device.py index d2565cf..64d5fc0 100644 --- a/belay/device.py +++ b/belay/device.py @@ -745,7 +745,7 @@ def setup( """ # noqa: D400 if f is None: return wraps_partial(Device.setup, autoinit=autoinit, **kwargs) # type: ignore[reportGeneralTypeIssues] - if signature(f).parameters: + if signature(f).parameters and autoinit: raise ValueError( f"Method {f} decorated with " '"@Device.setup(autoinit=True)" ' diff --git a/tests/integration/test_classes.py b/tests/integration/test_classes.py index 3f000bf..992d3a8 100644 --- a/tests/integration/test_classes.py +++ b/tests/integration/test_classes.py @@ -33,6 +33,30 @@ def get_times_bar(val): assert 84 == device.get_times_bar(2) +def test_classes_setup_arguments(emulate_command): + class MyDevice(Device, skip=True): + @Device.setup + def setup1(baz=1): + foo = 11 # noqa: F841 + bar = 41 # noqa: F841 + + @Device.setup() + def setup2(baz=2): + foo = 12 # noqa: F841 + bar = 42 # noqa: F841 + + with MyDevice(emulate_command) as device: + device.setup1(baz=111) + assert device("foo") == 11 + assert device("bar") == 41 + assert device("baz") == 111 + + device.setup2(baz=222) + assert device("foo") == 12 + assert device("bar") == 42 + assert device("baz") == 222 + + def test_classes_setup_autoinit(emulate_command): class MyDevice(Device, skip=True): @Device.setup(autoinit=True) diff --git a/tests/integration/test_function_decorators.py b/tests/integration/test_function_decorators.py index d1dd557..28a89e7 100644 --- a/tests/integration/test_function_decorators.py +++ b/tests/integration/test_function_decorators.py @@ -1,6 +1,7 @@ import pytest import belay +from belay import Device def test_setup_basic(emulated_device): @@ -77,3 +78,24 @@ def foo(): emulated_device.close() mock_teardown.assert_called_once() + + +def test_classdecorator_setup(): + @Device.setup + def foo1(): + pass + + @Device.setup() + def foo2(): + pass + + @Device.setup(autoinit=True) + def foo3(): + pass + + with pytest.raises(ValueError): + # Provided an arg with autoinit=True is not allowed. + + @Device.setup(autoinit=True) + def foo(arg1=1): + pass