Source code for schrodinger.test.mock_ld_models

from schrodinger.application.livedesign import login
from schrodinger.application.livedesign.live_report_widget import LiveReportType

PROJ_IDS = [login.GLOBAL_PROJECT_ID, '1', '2']
DEFAULT_PROJ_ID = PROJ_IDS[1]


class MockNameAndID:

    def __init__(self, name=None, id_=None):
        super().__init__()
        self._name = name
        self.id = id_

    @property
    def name(self):
        return self._name


class MockLDProject(MockNameAndID):

    def __init__(self, name=None, id_=DEFAULT_PROJ_ID):
        super().__init__(name, id_)


class MockFolder(MockNameAndID):

    def __init__(self, name=None, id_=None, project_id=DEFAULT_PROJ_ID):
        super().__init__(name, id_)
        self.project_id = project_id


class MockLiveReport:

    def __init__(self,
                 title=None,
                 template=None,
                 owner=None,
                 tags=None,
                 id_=None,
                 project_id=DEFAULT_PROJ_ID,
                 lr_type=LiveReportType.COMPOUND):
        super().__init__()
        self.title = title
        self.template = template
        self.owner = owner
        if tags is None:
            tags = []
        self.tags = tags
        self.id = id_
        self.project_id = project_id
        self.type = lr_type


class MockLDModels:
    """
    Mock the `ld_models` module; treat classes as methods that return new mock
    instances of the corresponding class.
    """

    def LiveReport(self,
                   title=None,
                   owner=None,
                   tags=None,
                   id_=None,
                   project_id=DEFAULT_PROJ_ID):
        return MockLiveReport(title=title,
                              owner=owner,
                              tags=tags,
                              id_=id_,
                              project_id=project_id)


class FreeformColumn:
    """
    A class meant to mock the relevant attributes of
    `ldclient.models.FreeformColumn`.
    """
    COLUMN_ATTACHMENT = 'ATTACHMENT'
    COLUMN_TEXT = 'TEXT'

    def __init__(self,
                 name,
                 description,
                 type=COLUMN_TEXT,
                 published=None,
                 project_id=None,
                 live_report_id=None):
        self._name = name
        self._description = description
        self._type = type
        self._published = published
        self._project_id = project_id
        self._live_report_id = live_report_id

    @property
    def name(self):
        return self._name

    @property
    def published(self):
        return self._published

    @property
    def project_id(self):
        return self._project_id

    @property
    def live_report_id(self):
        return self._live_report_id

    @property
    def id(self):
        return id(self)

    def __eq__(self, obj):
        if not isinstance(obj, FreeformColumn):
            return False
        return (self.published == obj.published and
                self.project_id == obj.project_id and
                self.live_report_id == obj.live_report_id and
                self.name == obj.name and self.id == obj.id)


class Column:
    """
    A class meant to mock the relevant attributes of
    `ldclient.models.Column`.
    """

    def __init__(self,
                 id=None,
                 name=None,
                 column_type=None,
                 value_type=None,
                 log_scale=None,
                 project_ids=None,
                 folder_name=None):
        self.id = id
        self.name = name
        self.column_type = column_type
        self.value_type = value_type
        self.log_scale = log_scale
        self.project_ids = project_ids
        self.folder_name = folder_name


class BaseModel:
    pass


class Model(BaseModel):
    """
    A class to mock the relevant attributes of `models.Model`.
    """

    def __init__(self,
                 name,
                 archived,
                 published,
                 user,
                 folder,
                 returns,
                 project_ids,
                 template_vars,
                 description,
                 parent=None,
                 commands=None,
                 batch_size=None,
                 command_type=None,
                 command_queue=None,
                 created_at=None,
                 updated_at=None,
                 id=None):
        self.id = id
        self.parent = parent
        self.name = name
        self.archived = archived
        self.published = published
        self.user = user
        self.created_at = created_at
        self.updated_at = updated_at
        self.folder = folder
        self.returns = returns
        self.template_vars = template_vars
        self.batch_size = batch_size
        self.description = description
        self.command_type = command_type
        self.command_queue = command_queue
        self.commands = [] if not commands else commands
        self.project_ids = project_ids or [login.GLOBAL_PROJECT_ID]


class ModelRecursive(BaseModel):
    """
    A class to mock the relevant attributes of `models.ModelRecursive`.
    """

    def __init__(self, tag, value=None):
        self.tag = tag
        self.value = value


class ModelTemplateVar:
    """
    Mock of `models.ModelTemplateVar`, "a binding of a variable in a command".
    """

    def __init__(self, tag, type, name, data=None, id=None):
        self.id = id
        self.tag = tag
        self.type = type
        self.name = name
        self.data = data


class AsyncServiceTaskFailedError(Exception):
    pass