Skip to content

Commit

Permalink
Make MockSignalBackend more performant
Browse files Browse the repository at this point in the history
  • Loading branch information
coretl committed May 17, 2024
1 parent ed2e434 commit 45cc406
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 27 deletions.
29 changes: 14 additions & 15 deletions src/ophyd_async/core/mock_signal_backend.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
from functools import cached_property
from typing import Optional, Type
from unittest.mock import MagicMock
from unittest.mock import Mock

from bluesky.protocols import Descriptor, Reading

Expand Down Expand Up @@ -36,51 +37,49 @@ def __init__(
else:
self.soft_backend = initial_backend

self.mock = MagicMock()

self.put_proceeds = asyncio.Event()
self.put_proceeds.set()

def source(self, name: str) -> str:
self.mock.source(name)
if self.initial_backend:
return f"mock+{self.initial_backend.source(name)}"
return f"mock+{name}"

async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
self.mock.connect(timeout=timeout)
pass

@cached_property
def put_mock(self) -> Mock:
return Mock(name="put")

@cached_property
def put_proceeds(self) -> asyncio.Event:
put_proceeds = asyncio.Event()
put_proceeds.set()
return put_proceeds

async def put(self, value: Optional[T], wait=True, timeout=None):
self.mock.put(value, wait=wait, timeout=timeout)
self.put_mock(value, wait=wait, timeout=timeout)
await self.soft_backend.put(value, wait=wait, timeout=timeout)

if wait:
await asyncio.wait_for(self.put_proceeds.wait(), timeout=timeout)

def set_value(self, value: T):
self.mock.set_value(value)
self.soft_backend.set_value(value)

async def get_descriptor(self, source: str) -> Descriptor:
self.mock.get_descriptor(source)
return await self.soft_backend.get_descriptor(source)

async def get_reading(self) -> Reading:
self.mock.get_reading()
return await self.soft_backend.get_reading()

async def get_value(self) -> T:
self.mock.get_value()
return await self.soft_backend.get_value()

async def get_setpoint(self) -> T:
"""For a soft signal, the setpoint and readback values are the same."""
self.mock.get_setpoint()
return await self.soft_backend.get_setpoint()

async def get_datakey(self, source: str) -> Descriptor:
return await self.soft_backend.get_datakey(source)

def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
self.mock.set_callback(callback)
self.soft_backend.set_callback(callback)
14 changes: 7 additions & 7 deletions src/ophyd_async/core/mock_signal_utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from contextlib import asynccontextmanager, contextmanager
from typing import Any, Callable, Generator, Iterable, Iterator, List
from unittest.mock import ANY
from unittest.mock import ANY, Mock

from ophyd_async.core.signal import Signal
from ophyd_async.core.utils import T
Expand Down Expand Up @@ -43,12 +43,12 @@ async def mock_puts_blocked(*signals: List[Signal]):

def assert_mock_put_called_with(signal: Signal, value: Any, wait=ANY, timeout=ANY):
backend = _get_mock_signal_backend(signal)
backend.mock.put.assert_called_with(value, wait=wait, timeout=timeout)
backend.put_mock.assert_called_with(value, wait=wait, timeout=timeout)


def reset_mock_put_calls(signal: Signal):
backend = _get_mock_signal_backend(signal)
backend.mock.put.reset_mock()
backend.put_mock.reset_mock()


class _SetValuesIterator:
Expand Down Expand Up @@ -122,9 +122,9 @@ def set_mock_values(


@contextmanager
def _unset_side_effect_cm(mock):
def _unset_side_effect_cm(put_mock: Mock):
yield
mock.put.side_effect = None
put_mock.side_effect = None


# linting isn't smart enought to realize @contextmanager will give use a
Expand All @@ -145,5 +145,5 @@ def callback_on_mock_put(
The callback to call when the backend is put to during the context.
"""
backend = _get_mock_signal_backend(signal)
backend.mock.put.side_effect = callback
return _unset_side_effect_cm(backend.mock)
backend.put_mock.side_effect = callback
return _unset_side_effect_cm(backend.put_mock)
6 changes: 1 addition & 5 deletions tests/core/test_mock_signal_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,7 @@ async def test_mock_signal_backend(connect_mock_mode):
assert await mock_signal._backend.get_value() == ""
await mock_signal._backend.put("test")
assert await mock_signal._backend.get_value() == "test"

mock_signal._backend.mock.get_value.assert_called_once

mock_signal._backend.mock["get_value"].assert_called_once
assert mock_signal._backend.mock.put.call_args_list == [
assert mock_signal._backend.put_mock.call_args_list == [
call("test", wait=True, timeout=None),
]

Expand Down

0 comments on commit 45cc406

Please sign in to comment.