Source code for pgmock.mocker


Contains the primary functionality for patching and mocking SQLAlchemy queries
import collections

import sqlalchemy.event as sqla_event
import sqlalchemy.orm.session as sqla_session

import pgmock.exceptions
import pgmock.render

Data = collections.namedtuple('Data', ['rows', 'cols'])

[docs]def data(rows=None, cols=None): """Creates patch data for a side effect. Args: rows (List[tuple], optional): A list of tuples of values to patch for each row. Each row must have the same length. If ``None``, defaults to an empty list. cols (List[str]): A list of columns. Returns: Data: A data object that can be used as input to a side effect of a patch, for example ``pgmock.patch(side_effect=[, cols=...)])`` """ return Data(rows=rows, cols=cols)
[docs]def mock(connectable, replace_new_patch_aliases=None): """Creates a mock selector that can be patched. This is intended to be used as a context manager with a given SQLAlchemy connectable (e.g. an engine, session, connection, etc). For example:: with pgmock.mock(engine) as mocker: mocker.patch(pgmock.table('my_table'), rows, cols) mocker.patch(pgmock.table('other_table'), rows, cols) # Run SQLAlchemy queries... # Assert the mocker was rendered with as many queries executed assert len(mocker.renderings) == num_expected_queries Any queries executed inside of the context manager will be patched by SQLAlchemy's ``before_cursor_execute`` event. Renderings of patched SQL can be obtained by examining the ``renderings`` property of the object, which is a list of tuples of the original and modified SQL of every query. If any of the patches cannot be matched during query execution, the relevant exceptions are raised. Specific patches can be applied to specific queries by using the ``side_effect`` argument of `pgmock.patch`. Args: connectable (SQLAlchemy connectable object): The connectable SQLAlchemy object (e.g engine, session, connection, etc) replace_new_patch_aliases (bool, optional): If ``True``, will replace any references to patch aliases when they differ from the original alias. If ``None``, uses the globally-configured value that defaults to ``True``. More information on this can be found at `pgmock.config.set_replace_new_patch_aliases` Returns: Mock: A chainable mock object that can be patched. Raises: Any error that can happen during rendering. """ return Mocker(connectable=connectable, replace_new_patch_aliases=replace_new_patch_aliases)
[docs]def patch(selector=None, rows=None, cols=None, side_effect=None): """Applies a patch to a selector. Args: selector (Selector, optional). A selector to patch inside of the relevant SQL. rows (List[tuple], optional): A list of tuples of values to patch for each row. Each row must have the same length. If ``None``, patching is ignored. cols (List[str]): A list of columns. If more columns are provided than the length of the rows, ``null`` values are filled in for the missing values. side_effect(List[]): A list of side effects. Side effects can only be provided when ``rows`` and ``cols`` are not provided. Each side effect is rendered on each subsequent rendering of the patch. Side effects must be instantiated with ```` and the arguments are ``rows`` and ``cols``. Note: providing ``None`` as a side effect will ignore the patch for the rendering. Returns: Mock: A chainable mock object that can be patched. Raises: `UnpatchableError`: When the selector cannot be patched Examples: Patch a table "schema.table_name" with values .. code-block:: python patch = pgmock.patch(pgmock.table('schema.table_name'), rows=[(1, 2), (3, 4)], cols=['a', 'b']) patched_query = pgmock.sql(sql_string, patch) Patch a table "schema.table_name" with a side effect while using SQLAlchemy .. code-block:: python with pgmock.mock(connetion) as mocker: mocker.patch(pgmock.table('schema.table_name'), side_effect=[ None,[(1, 2), (3, 4)], ['a', 'b']) ]) # Do no patching on the first execution of the SQLAlchemy # connection since the side effect returns ``None`` the # first time connection.execute(...) # Now apply the patch the second time connection.execute(...) """ return Mocker().patch(selector=selector, rows=rows, cols=cols, side_effect=side_effect)
class SideEffect: """A side effect for a mock Similar to traditional mock side effects, ``pgmock`` side effects are callables that return different results for each call. Currently ``pgmock`` only supports taking an iterable as a side effect's input. Taking another callable is intended to be supported in future versions """ def __init__(self, side_effect): """Initializes the side effect and ensures its the proper type""" #: The index into the side effect self.iterable_idx = 0 self.side_effect = side_effect if not isinstance(side_effect, (list, tuple)): raise TypeError('Side effects must be iterable') for se in self.side_effect: if not isinstance(se, Data) and se is not None: # While users could technically pass tuples, it's easy to get the ordering mixed up # with rows and cols raise TypeError('Side effect values must be instantiated with') def __call__(self): if self.iterable_idx >= len(self.side_effect): # Raise custom error message pgmock.exceptions.throw(pgmock.exceptions.SideEffectExhaustedError, 'Side effect of length {} has been exhausted'.format( len(self.side_effect))) self.iterable_idx += 1 return self.side_effect[self.iterable_idx - 1] class Mocker(pgmock.render.Renderable): """A renderable for patching expressions and tables When used as a context manager, the user must supply a SQLAlchemy connectable object (eg. Engine, Session, Connection, etc) to the constructor:: with Mock(connectable=connectable) as mocker: ... When used as a context manager, the ``before_cursor_execute`` event is tracked and queries are patched when issued. """ def __init__(self, renderable=None, connectable=None, replace_new_patch_aliases=None): super().__init__(renderable=renderable) if replace_new_patch_aliases is None: replace_new_patch_aliases = pgmock.config.get_replace_new_patch_aliases() self.replace_new_patch_aliases = replace_new_patch_aliases self._replace_new_patch_aliases_config = None self._query_hook = None # If the caller gives us a session, take the underlying connection or # engine instead. if isinstance(connectable, sqla_session.Session): # pragma: no cover if connectable.bind is None: raise TypeError("Can't use unbound `Session` object for mocking.") connectable = connectable.bind self._connectable = connectable @pgmock.render.Renderable.chainable_render_method def patch(self, selector=None, rows=None, cols=None, side_effect=None): """Returns a chainable render method and ensures a side effect object is passed to the renderer""" # pylint: disable=no-self-use if selector and not isinstance(selector, pgmock.render.Renderable): # pragma: no cover raise TypeError( 'Must provide a selector to patch. Type provided = "{}"'.format(type(selector))) if side_effect and not isinstance(side_effect, SideEffect): side_effect = SideEffect(side_effect) return pgmock.render.RenderMethod(name='patch', args=[], kwargs={ 'selector': selector, 'rows': rows, 'cols': cols, 'side_effect': side_effect }) def start(self): """Starts the ``before_cursor_execute`` event listener""" if not self._connectable: raise pgmock.exceptions.throw(pgmock.exceptions.NoConnectableError, 'Must provide a connectable when using context manager') self.stop() # pylint: disable=unused-argument @sqla_event.listens_for(self._connectable, 'before_cursor_execute', retval=True) def _hook(conn, cursor, statement, parameters, context, executemany): """Query hook to apply patches""" statement = self.render(statement).sql_view return statement, parameters # Set up our query modifier to listen for execution events self._query_hook = _hook sqla_event.listen(self._connectable, 'before_cursor_execute', _hook) # Set the configuration for replacing new patch aliases self._replace_new_patch_aliases_config = pgmock.config.set_replace_new_patch_aliases( self.replace_new_patch_aliases) self._replace_new_patch_aliases_config.__enter__() return self def stop(self): """Stops the ``before_cursor_execute`` event listener""" if self._query_hook: sqla_event.remove(self._connectable, 'before_cursor_execute', self._query_hook) self._query_hook = None if self._replace_new_patch_aliases_config: self._replace_new_patch_aliases_config.__exit__(None, None, None) self._replace_new_patch_aliases_config = None def __enter__(self): return self.start() def __exit__(self, *error_args): self.stop()