Source code for schrodinger.test.project_mock

"""
This module has classes that can be used to mock out project related
functionality that cannot be accessed from outside of Maestro for unit tests.

An example of how to use would be to have a setUp function in a
`unittest.TestCase` derived class that does the following:

class MyTestClass(unittest.TestCase):

    def setUp(self):
        self.mock_mae = MagicMock()
        mymodule.maestro = self.mock_mae
        self.pt = MockProjectTable('project_sts.maegz')
        self.mock_mae.project_table_get = MagicMock(return_value=self.pt)

Copyright Schrodinger, LLC. All rights reserved.
"""

from schrodinger import structure


class MockProjectRow:
    """
    Standin for the `schrodinger.project.ProjectRow` class for tests.
    """

    def __init__(self, eid, st):
        """
        :param st: Structure associated with this row
        :type st: `schrodinger.structure.Structure`
        """
        self.entry_id = int(eid)
        self._st = st
        self.title = st.title
        self.in_workspace = False

    def getStructure(self, props=True, copy=True, workspace_sync=True):
        """
        Return the structure for this row.

        :param props: Ignored (properties are always included)
        :type props: bool

        :param copy: Whether the return value should be
            the structure object itself or a copy thereof
        :type copy: bool

        :param workspace_sync: Ignored
        :param workspace_sync: bool

        :return: The Structure for this row
        :rtype: `structure.Structure`
        """
        if copy:
            return self._st.copy()
        else:
            return self._st

    def setStructure(self, st):
        """
        Set the row structure to the specified structure object.

        :param st: New structure for this row
        :type st: `structure.Structure`
        """
        self._st = st

    def includeOnly(self):
        """
        In tests we do nothing for workspace inclusion/exclusion.
        """
        pass

    def __getitem__(self, propname):
        return self._st.property.get(propname)


class MockProjectTable:
    """
    Stands in for the `schrodinger.project.Project` class for the tests.
    """

    def __init__(self, st_fname=None, sts=None):
        """
        :param st_fname: Name of structure file to load project row structures from.
        :type st_fname: str

        :param sts: Structures to load
        :type sts: list of `structrue`
        """
        self._rows = []
        if st_fname:
            self.importStructureFile(st_fname)
        elif sts:
            for st in sts:
                self.importStructure(st)

    def _nextEntryId(self):
        if not self._rows:
            return 1
        return max(int(row.entry_id) for row in self._rows) + 1

    def importStructureFile(self,
                            filename,
                            wsreplace=True,
                            creategroups="multiple"):
        sr = structure.StructureReader(filename)
        start_eid = self._nextEntryId()
        for eid, st in enumerate(sr, start=start_eid):
            row = MockProjectRow(eid, st)
            self._rows.append(row)
        if wsreplace:
            self.includeRows([start_eid])

    def importStructure(self, st, name=None, wsreplace=False, copy=True):
        eid = self._nextEntryId()
        row = MockProjectRow(eid, st)
        self._rows.append(row)
        if wsreplace:
            self.includeRows([eid])
        return row

    def getRow(self, entry_id):
        """
        Return the row at the specified eid index or None if not found.

        :param entry_id: 'Entry ID' index for the row to return
        :type entry_id: int or str

        :return: The row of the specified entry if found or None if not
        :rtype: `FakeProjectRow` or None
        """
        entry_id = int(entry_id)
        for row in self._rows:
            if row.entry_id == entry_id:
                return row
        return None

    def __getitem__(self, eid):
        row = self.getRow(eid)
        if row is None:
            raise KeyError
        return row

    def includeRows(self, entry_ids, exclude_others=True, autofit=True):
        """
        Mock project row inclusion  in the workspace. The arguments here are
        added to mimic project.Project.includeRows.

        :param entry_ids: Entry IDs to include.
        :type entry_ids: list[int]

        :param exclude_others: Whether to exclude previous included entries.
        :type exclude_others: bool

        :param autofit: Used only to match API call for
                project.Project.includeRows, any value passed will be
                ignored as this relates to Maestro workspace.
        :type autofit: bool
        """
        for row in self._rows:
            if row.entry_id in entry_ids:
                row.in_workspace = True
            elif exclude_others:
                row.in_workspace = False

    included_rows = property(lambda self:
                             (row for row in self._rows if row.in_workspace))
    # For now, all rows are always selected:
    selected_rows = property(lambda self: self._rows)
    all_rows = property(lambda self: self._rows)