liminal-orm 3.2.1__py3-none-any.whl → 4.0.1__py3-none-any.whl

liminal/orm/base_results_model.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import inspect
+import logging
+from types import FunctionType
+from typing import Any, Generic, TypeVar  # noqa: UP035
+
+import pandas as pd  # type: ignore
+from sqlalchemy import DATETIME, Boolean, ForeignKey, String
+from sqlalchemy import Column as SqlColumn
+from sqlalchemy.orm import Query, RelationshipProperty, Session, relationship
+from sqlalchemy.orm.decl_api import declared_attr
+
+from liminal.base.base_validation_filters import BaseValidatorFilters
+from liminal.orm.base import Base
+from liminal.orm.base_tables.user import User
+from liminal.orm.results_schema_properties import ResultsSchemaProperties
+from liminal.validation import BenchlingValidatorReport
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+class BaseResultsModel(Generic[T], Base):
+    """Base class for all results models. Defines all common columns for results tables in Postgres Benchling warehouse."""
+
+    __abstract__ = True
+    __schema_properties__: ResultsSchemaProperties
+
+    @declared_attr
+    def creator_id(cls) -> SqlColumn:
+        return SqlColumn(
+            "creator_id$", String, ForeignKey("user$raw.id"), nullable=True
+        )
+
+    @declared_attr
+    def creator(cls) -> RelationshipProperty:
+        return relationship("User", foreign_keys=[cls.creator_id])
+
+    id = SqlColumn("id", String, nullable=True, primary_key=True)
+    archived = SqlColumn("archived$", Boolean, nullable=True)
+    archive_purpose = SqlColumn("archive_purpose$", String, nullable=True)
+    created_at = SqlColumn("created_at$", DATETIME, nullable=True)
+    entry_id = SqlColumn("entry_id$", String, nullable=True)
+    modified_at = SqlColumn("modified_at$", DATETIME, nullable=True)
+    run_id = SqlColumn("run_id$", String, nullable=True)
+    v3_id = SqlColumn("v3_id$", String, nullable=True)
+
+    def __init_subclass__(cls, **kwargs: Any):
+        super().__init_subclass__(**kwargs)
+        warehouse_name = cls.__schema_properties__.warehouse_name
+        cls.__tablename__ = warehouse_name + "$raw"
+
+    @classmethod
+    def apply_base_filters(
+        cls,
+        query: Query,
+        filter_archived: bool = True,
+        filter_unregistered: bool = True,
+        base_filters: BaseValidatorFilters | None = None,
+    ) -> Query:
+        """Applies the base model filters to the given query."""
+        if filter_archived:
+            query = query.filter(cls.archived.is_(False))
+        if filter_unregistered:
+            if hasattr(cls, "is_registered"):
+                query = query.filter(cls.is_registered.is_(True))
+
+        if base_filters is None:
+            return query
+        if base_filters.created_date_start:
+            query = query.filter(cls.created_at >= base_filters.created_date_start)
+        if base_filters.created_date_end:
+            query = query.filter(cls.created_at <= base_filters.created_date_end)
+        if base_filters.updated_date_start:
+            query = query.filter(cls.modified_at >= base_filters.updated_date_start)
+        if base_filters.updated_date_end:
+            query = query.filter(cls.modified_at <= base_filters.updated_date_end)
+        if base_filters.entity_ids:
+            query = query.filter(cls.v3_id.in_(base_filters.entity_ids))
+        if base_filters.creator_full_names:
+            query = query.filter(User.name.in_(base_filters.creator_full_names))
+        return query
+
+    @classmethod
+    def all(cls, session: Session) -> list[T]:
+        """Uses the get_query method to retrieve all results schema rows from the database.
+
+        Parameters
+        ----------
+        session : Session
+            Benchling database session.
+
+        Returns
+        -------
+        list[T]
+            List of all results schema rows from the database.
+        """
+        return cls.query(session).all()
+
+    @classmethod
+    def df(cls, session: Session) -> pd.DataFrame:
+        """Uses the get_query method to retrieve all results schema rows from the database.
+
+        Parameters
+        ----------
+        session : Session
+            Benchling database session.
+
+        Returns
+        -------
+        pd.DataFrame
+            A pandas dataframe of all results schema rows from the database.
+        """
+        query = cls.query(session)
+        return pd.read_sql(query.statement, session.connection())
+
+    @classmethod
+    def query(cls, session: Session) -> Query:
+        """Abstract method that users can override to define a specific query
+        to retrieve results schema rows from the database and cover any distinct relationships.
+
+        Parameters
+        ----------
+        session : Session
+            Benchling database session.
+
+        Returns
+        -------
+        Query
+            sqlalchemy query to retrieve results schema rows from the database.
+        """
+        return session.query(cls)
+
+    @classmethod
+    def get_validators(cls) -> list[FunctionType]:
+        """Returns a list of all validators defined on the class. Validators are functions that are decorated with @validator."""
+        validators = []
+        for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
+            if hasattr(method, "_is_liminal_validator"):
+                validators.append(method)
+        return validators
+
+    @classmethod
+    def validate(
+        cls,
+        session: Session,
+        base_filters: BaseValidatorFilters | None = None,
+        only_invalid: bool = False,
+    ) -> list[BenchlingValidatorReport]:
+        """Runs all validators for all results schema rows returned from the query and returns a list of reports.
+        This returns a report for each results schema row, validator pair, regardless of whether the validation passed or failed.
+
+        Parameters
+        ----------
+        session : Session
+            Benchling database session.
+        base_filters: BaseValidatorFilters
+            Filters to apply to the query.
+        only_invalid: bool
+            If True, only returns reports for entities that failed validation.
+
+        Returns
+        -------
+        list[BenchlingValidatorReport]
+            List of reports from running all validators on all results schema rows returned from the query.
+        """
+        results: list[BenchlingValidatorReport] = []
+        table: list[T] = cls.apply_base_filters(
+            cls.query(session), base_filters=base_filters
+        ).all()
+        logger.info(
+            f"Validating {len(table)} results schema rows for {cls.__name__}..."
+        )
+        validator_functions = cls.get_validators()
+        for entity in table:
+            for validator_func in validator_functions:
+                report: BenchlingValidatorReport = validator_func(entity)
+                if only_invalid and report.valid:
+                    continue
+                results.append(report)
+        return results
+
+    @classmethod
+    def validate_to_df(
+        cls,
+        session: Session,
+        base_filters: BaseValidatorFilters | None = None,
+        only_invalid: bool = False,
+    ) -> pd.DataFrame:
+        """Runs all validators for all results schema rows returned from the query and returns reports as a pandas dataframe.
+
+        Parameters
+        ----------
+        session : Session
+            Benchling database session.
+        base_filters: BaseValidatorFilters
+            Filters to apply to the query.
+
+        Returns
+        -------
+        pd.Dataframe
+            Dataframe of reports from running all validators on all results schema rows returned from the query.
+        """
+        results = cls.validate(session, base_filters, only_invalid)
+        return pd.DataFrame([r.model_dump() for r in results])

liminal/orm/relationship.py

@@ -36,7 +36,7 @@ def single_relationship(
 
 
 def multi_relationship(*args: Any, **kwargs: Any) -> RelationshipProperty:
-    """Wrapper for generating a multi-relationship. Supporting the usage of a deprecated signature until v4 release."""
+    """Wrapper for generating a multi-relationship. Supporting the usage of a deprecated signature until v5 release."""
     if len(args) == 2 and isinstance(args[1], Column):
         return multi_relationship_v2(*args, **kwargs)
     else:
@@ -64,7 +64,7 @@ def multi_relationship_deprecated(
     SQLAlchemy RelationshipProperty
     """
     warnings.warn(
-        "This version of multi_relationship is deprecated. New function signature is multi_relationship(target_class_name: str, entity_link_field: Column). Support for this signature will end with the v4 release.",
+        "This version of multi_relationship is deprecated. New function signature is multi_relationship(target_class_name: str, entity_link_field: Column). Support for this signature will end with the v5 release.",
         FutureWarning,
         stacklevel=2,
     )

liminal/orm/results_schema_properties.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class ResultsSchemaProperties(BaseModel):
+    """
+    This class is the validated class that is public facing, and represents the properties of a results schema properties.
+
+    Parameters
+    ----------
+    name : str
+        The name of the schema.
+    warehouse_name : str
+       The sql table name of the schema in the benchling warehouse.
+    """
+
+    name: str
+    warehouse_name: str
+
+    def __repr__(self) -> str:
+        """Generates a string representation of the class so that it can be executed."""
+        return f"{self.__class__.__name__}({', '.join([f'{k}={v.__repr__()}' for k, v in self.model_dump().items()])})"

liminal/orm/schema_properties.py

@@ -94,3 +94,7 @@ class SchemaProperties(BaseSchemaProperties):
     def set_archived(self, value: bool) -> SchemaProperties:
         self._archived = value
         return self
+
+    def set_warehouse_name(self, value: str) -> SchemaProperties:
+        self.warehouse_name = value
+        return self

liminal/results_schemas/generate_files.py

@@ -0,0 +1,197 @@
+import shutil
+from pathlib import Path
+
+from rich import print
+
+from liminal.base.base_dropdown import BaseDropdown
+from liminal.connection import BenchlingService
+from liminal.dropdowns.utils import get_benchling_dropdowns_dict
+from liminal.entity_schemas.utils import get_converted_tag_schemas
+from liminal.enums import BenchlingFieldType
+from liminal.mappers import convert_benchling_type_to_python_type
+from liminal.orm.base_model import BaseModel
+from liminal.results_schemas.utils import get_converted_results_schemas
+from liminal.utils import to_pascal_case, to_snake_case
+
+TAB = "    "
+
+
+def generate_all_results_schema_files(
+    benchling_service: BenchlingService, write_path: Path, overwrite: bool = False
+) -> None:
+    """Generate all results schema files from your Benchling tenant and writes to the given results_schemas/ path.
+    This is used to initialize your code for Liminal and transfer the information from your Benchling tenant to your local codebase.
+
+    Parameters
+    ----------
+    benchling_service : BenchlingService
+        The Benchling service object that is connected to a specified Benchling tenant.
+    write_path : Path
+        The path to write the generated files to. results_schemas/ directory will be created within this path.
+    overwrite : bool
+        Whether to overwrite existing the existing results_schemas/ directory.
+    """
+    write_path = write_path / "results_schemas"
+    if write_path.exists() and overwrite:
+        shutil.rmtree(write_path)
+        print(f"[dim]Removed directory: {write_path}")
+    if not write_path.exists():
+        write_path.mkdir(parents=True, exist_ok=True)
+        print(f"[green]Created directory: {write_path}")
+
+    results_schemas = get_converted_results_schemas(benchling_service)
+    entity_schemas_wh_name_to_classname = _get_entity_schemas_wh_name_to_classname(
+        benchling_service
+    )
+    dropdown_name_to_classname_map = _get_dropdown_name_to_classname_map(
+        benchling_service
+    )
+    init_file_imports = []
+
+    for schema_properties, field_properties_dict in results_schemas:
+        has_date = False
+        file_name = to_snake_case(schema_properties.warehouse_name) + ".py"
+        schema_name = to_pascal_case(schema_properties.warehouse_name)
+        init_file_imports.append(
+            f"from .{to_snake_case(schema_properties.warehouse_name)} import {schema_name}"
+        )
+        import_strings = [
+            "from sqlalchemy import Column as SqlColumn",
+            "from liminal.orm.base_results_model import BaseResultsModel",
+            "from liminal.orm.results_schema_properties import ResultsSchemaProperties",
+            "from liminal.orm.column import Column",
+            "from liminal.enums import BenchlingFieldType",
+        ]
+        init_strings = [f"{TAB}def __init__(", f"{TAB}self,"]
+        column_strings = []
+        dropdowns = []
+        relationship_strings = []
+        for col_name, col in field_properties_dict.items():
+            column_props = col.column_dump()
+            dropdown_classname = None
+            if col.dropdown_link:
+                dropdown_classname = dropdown_name_to_classname_map[col.dropdown_link]
+                dropdowns.append(dropdown_classname)
+                column_props["dropdown_link"] = dropdown_classname
+            column_props_string = ""
+            for k, v in column_props.items():
+                if k == "dropdown_link":
+                    column_props_string += f"""dropdown={v},"""
+                else:
+                    column_props_string += f"""{k}={v.__repr__()},"""
+            column_string = f"""{TAB}{col_name}: SqlColumn = Column({column_props_string.rstrip(',')})"""
+            column_strings.append(column_string)
+            if col.required and col.type:
+                init_strings.append(
+                    f"""{TAB}{col_name}: {convert_benchling_type_to_python_type(col.type).__name__},"""
+                )
+
+            if (
+                col.type == BenchlingFieldType.DATE
+                or col.type == BenchlingFieldType.DATETIME
+            ):
+                if not has_date:
+                    import_strings.append("from datetime import datetime")
+            if (
+                col.type in BenchlingFieldType.get_entity_link_types()
+                and col.entity_link is not None
+            ):
+                if not col.is_multi:
+                    relationship_strings.append(
+                        f"""{TAB}{col_name}_entity = single_relationship("{entity_schemas_wh_name_to_classname[col.entity_link]}", {col_name})"""
+                    )
+                    import_strings.append(
+                        "from liminal.orm.relationship import single_relationship"
+                    )
+                else:
+                    relationship_strings.append(
+                        f"""{TAB}{col_name}_entities = multi_relationship("{entity_schemas_wh_name_to_classname[col.entity_link]}", {col_name})"""
+                    )
+                    import_strings.append(
+                        "from liminal.orm.relationship import multi_relationship"
+                    )
+        for col_name, col in field_properties_dict.items():
+            if not col.required and col.type:
+                init_strings.append(
+                    f"""{TAB}{col_name}: {convert_benchling_type_to_python_type(col.type).__name__} | None = None,"""
+                )
+        init_strings.append("):")
+        for col_name in field_properties_dict.keys():
+            init_strings.append(f"{TAB}self.{col_name} = {col_name}")
+        if len(dropdowns) > 0:
+            import_strings.append(f"from ..dropdowns import {', '.join(dropdowns)}")
+        for col_name, col in field_properties_dict.items():
+            if col.dropdown_link:
+                init_strings.append(
+                    TAB
+                    + dropdown_name_to_classname_map[col.dropdown_link]
+                    + f".validate({col_name})"
+                )
+
+        import_string = "\n".join(list(set(import_strings)))
+        columns_string = "\n".join(column_strings)
+        relationship_string = "\n".join(relationship_strings)
+        init_string = (
+            f"\n{TAB}".join(init_strings) if len(field_properties_dict) > 0 else ""
+        )
+        schema_content = f"""{import_string}
+
+
+class {schema_name}(BaseResultsModel):
+    __schema_properties__ = {schema_properties.__repr__()}
+
+{columns_string}
+
+{relationship_string}
+
+{init_string}
+"""
+
+        with open(write_path / file_name, "w") as file:
+            file.write(schema_content)
+
+    with open(write_path / "__init__.py", "w") as file:
+        file.write("\n".join(init_file_imports))
+    print(
+        f"[green]Generated {write_path / '__init__.py'} with {len(results_schemas)} entity schema imports."
+    )
+
+
+def _get_dropdown_name_to_classname_map(
+    benchling_service: BenchlingService,
+) -> dict[str, str]:
+    """Gets the dropdown name to classname map.
+    If there are dropdowns imported, use BenchlingDropdown.get_all_subclasses()
+    Otherwise, it will query for Benchling dropdowns and use those.
+    """
+    if len(BaseDropdown.get_all_subclasses()) > 0:
+        return {
+            dropdown.__benchling_name__: dropdown.__name__
+            for dropdown in BaseDropdown.get_all_subclasses()
+        }
+    benchling_dropdowns = get_benchling_dropdowns_dict(benchling_service)
+    if len(benchling_dropdowns) > 0:
+        raise Exception(
+            "No dropdowns found locally. Please ensure your env.py file imports your dropdown classes or generate dropdowns from your Benchling tenant first."
+        )
+    return {}
+
+
+def _get_entity_schemas_wh_name_to_classname(
+    benchling_service: BenchlingService,
+) -> dict[str, str]:
+    """Gets the entity schema warehouse name to classname map.
+    If there are entity schemas imported, use BenchlingEntitySchema.get_all_subclasses()
+    Otherwise, it will query for Benchling entity schemas and use those.
+    """
+    if len(BaseModel.get_all_subclasses()) > 0:
+        return {
+            s.__schema_properties__.warehouse_name: s._sa_class_manager.class_.__name__
+            for s in BaseModel.get_all_subclasses()
+        }
+    tag_schemas = get_converted_tag_schemas(benchling_service)
+    if len(tag_schemas) > 0:
+        raise Exception(
+            "No entity schemas found locally. Please ensure your env.py file imports your entity schema classes or generate entity schemas from your Benchling tenant first."
+        )
+    return {}

liminal/results_schemas/models/results_schema_model.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any
+
+import requests
+from pydantic import BaseModel
+
+from liminal.connection.benchling_service import BenchlingService
+from liminal.entity_schemas.tag_schema_models import TagSchemaFieldModel
+
+
+class ResultsSchemaModel(BaseModel):
+    """A pydantic model to define a results schema, which is used when querying for results schemas from Benchling's internal API."""
+
+    allFields: list[TagSchemaFieldModel]
+    archiveRecord: dict[str, str] | None
+    derivedParent: Any | None
+    fields: list[TagSchemaFieldModel]
+    id: str
+    name: str | None
+    organization: Any | None
+    permissions: dict[str, bool] | None
+    prefix: str | None
+    publishedDataTableColumns: Any | None
+    requestTaskSchemaIds: list[Any] | None
+    requestTemplateIds: list[Any] | None
+    sampleGroupSchema: Any | None
+    schemaType: str
+    sqlIdentifier: str | None
+
+    @classmethod
+    def get_all_json(
+        cls,
+        benchling_service: BenchlingService,
+    ) -> list[dict[str, Any]]:
+        """This function gets all results schemas from Benchling's internal API, returning the raw JSON data.
+
+        Parameters
+        ----------
+        benchling_service : BenchlingService
+            The benchling service to use to get the results schemas.
+
+        Returns
+        -------
+        list[dict[str, Any]]
+            A list of results schemas, in their raw JSON format.
+        """
+
+        with requests.Session() as session:
+            response = session.get(
+                f"https://{benchling_service.benchling_tenant}.benchling.com/1/api/result-schemas",
+                headers=benchling_service.custom_post_headers,
+                cookies=benchling_service.custom_post_cookies,
+            )
+        if not response.ok:
+            raise Exception("Failed to get result schemas.")
+        return response.json()["data"]
+
+    @classmethod
+    def get_all(
+        cls,
+        benchling_service: BenchlingService,
+        wh_schema_names: set[str] | None = None,
+    ) -> list[ResultsSchemaModel]:
+        """This function gets all results schemas from Benchling's internal API.
+        If a list of warehouse names is provided, the function will only return the results schemas with the given warehouse names.
+
+        Parameters
+        ----------
+        benchling_service : BenchlingService
+            The benchling service to use to get the results schemas.
+        wh_schema_names : set[str] | None, optional
+            The set of warehouse names to filter the results schemas by. If not provided, all results schemas will be returned.
+
+        Returns
+        -------
+        list[ResultsSchemaModel]
+            A list of results schema models.
+        """
+        schemas_data = cls.get_all_json(benchling_service)
+        filtered_schemas: list[ResultsSchemaModel] = []
+        if wh_schema_names:
+            for schema in schemas_data:
+                if schema["sqlIdentifier"] in wh_schema_names:
+                    filtered_schemas.append(cls.model_validate(schema))
+                if len(filtered_schemas) == len(wh_schema_names):
+                    break
+        else:
+            for schema in schemas_data:
+                try:
+                    filtered_schemas.append(cls.model_validate(schema))
+                except Exception as e:
+                    print(f"Error validating schema {schema['sqlIdentifier']}: {e}")
+        return filtered_schemas
+
+    @classmethod
+    def get_one(
+        cls,
+        benchling_service: BenchlingService,
+        wh_schema_name: str,
+        schemas_data: list[dict[str, Any]] | None = None,
+    ) -> ResultsSchemaModel:
+        """This function gets a singular results schema, and raises an error if a schema with the given warehouse name is not found.
+
+        Parameters
+        ----------
+        benchling_service : BenchlingService
+            The benchling service to use to get the results schema.
+        wh_schema_name : str
+            The warehouse name of the results schema to search for.
+        schemas_data : list[dict[str, Any]] | None
+            The list of results schemas to search through, to avoid making extra API calls. If not provided, the function will get all results schemas from Benchling.
+
+        Returns
+        -------
+        ResultsSchemaModel
+            The corresponding results schema model.
+        """
+        if schemas_data is None:
+            schemas_data = cls.get_all_json(benchling_service)
+        schema = next(
+            (
+                schema
+                for schema in schemas_data
+                if schema["sqlIdentifier"] == wh_schema_name
+                and schema["registryId"] == benchling_service.registry_id
+            ),
+            None,
+        )
+        if schema is None:
+            raise ValueError(
+                f"Schema {wh_schema_name} not found in Benchling {benchling_service.benchling_tenant}."
+            )
+        return cls.model_validate(schema)
+
+    @classmethod
+    @lru_cache(maxsize=100)
+    def get_one_cached(
+        cls,
+        benchling_service: BenchlingService,
+        wh_schema_name: str,
+    ) -> ResultsSchemaModel:
+        """This function gets a singular results schema from Benchling and caches it."""
+        return cls.get_one(benchling_service, wh_schema_name)

liminal/results_schemas/utils.py

@@ -0,0 +1,48 @@
+from benchling_api_client.v2.stable.models.assay_result_schema import AssayResultSchema
+
+from liminal.base.properties.base_field_properties import BaseFieldProperties
+from liminal.connection import BenchlingService
+from liminal.dropdowns.utils import get_benchling_dropdown_id_name_map
+from liminal.entity_schemas.utils import convert_tag_schema_field_to_field_properties
+from liminal.orm.results_schema_properties import ResultsSchemaProperties
+from liminal.results_schemas.models.results_schema_model import ResultsSchemaModel
+from liminal.unit_dictionary.utils import get_unit_id_to_name_map
+
+
+def get_converted_results_schemas(
+    benchling_service: BenchlingService,
+) -> list[tuple[ResultsSchemaProperties, dict[str, BaseFieldProperties]]]:
+    """This functions gets all Results Schema schemas from Benchling and converts them to our internal representation of a schema and its fields.
+    It parses the Results Schema and creates ResultsSchemaProperties and a list of FieldProperties for each field in the schema.
+    """
+    results_schemas = ResultsSchemaModel.get_all(benchling_service)
+    dropdowns_map = get_benchling_dropdown_id_name_map(benchling_service)
+    unit_id_to_name_map = get_unit_id_to_name_map(benchling_service)
+    results_schemas_list = []
+    for schema in results_schemas:
+        schema_properties = ResultsSchemaProperties(
+            name=schema.name,
+            warehouse_name=schema.sqlIdentifier,
+        )
+        field_properties_dict = {}
+        for field in schema.fields:
+            field_properties_dict[field.systemName] = (
+                convert_tag_schema_field_to_field_properties(
+                    field, dropdowns_map, unit_id_to_name_map
+                )
+            )
+        results_schemas_list.append((schema_properties, field_properties_dict))
+    return results_schemas_list
+
+
+def get_results_schemas_dict(
+    benchling_service: BenchlingService,
+) -> dict[str, AssayResultSchema]:
+    """This function gets all Results Schema schemas using the Benchling API and returns a dictionary of the schemas by their system name."""
+    flattened_schemas = [
+        s
+        for schemas in list(benchling_service.schemas.list_assay_result_schemas())
+        for s in schemas
+    ]
+    schemas_dict = {s.system_name: s for s in flattened_schemas}
+    return schemas_dict

liminal/tests/test_entity_schema_compare.py

@@ -41,23 +41,22 @@ class TestCompareEntitySchemas:
 
             mock_get_benchling_entity_schemas.assert_called_once()
             mock_get_all_subclasses.assert_called()
-            assert len(invalid_models["mock_entity_two_wh"]) == 2
+            assert len(invalid_models["mock_entity_two"]) == 2
             assert isinstance(
-                invalid_models["mock_entity_two_wh"][0].op, CreateEntitySchema
+                invalid_models["mock_entity_two"][0].op, CreateEntitySchema
             )
             assert (
-                invalid_models["mock_entity_two_wh"][0].op.schema_properties.name
+                invalid_models["mock_entity_two"][0].op.schema_properties.name
                 == "Mock Entity Two"
             )
             assert [
-                f.warehouse_name
-                for f in invalid_models["mock_entity_two_wh"][0].op.fields
+                f.warehouse_name for f in invalid_models["mock_entity_two"][0].op.fields
             ] == ["parent_link_field"]
             assert isinstance(
-                invalid_models["mock_entity_two_wh"][1].op, UpdateEntitySchema
+                invalid_models["mock_entity_two"][1].op, UpdateEntitySchema
             )
             assert (
-                invalid_models["mock_entity_two_wh"][1].op.update_props.warehouse_name
+                invalid_models["mock_entity_two"][1].op.update_props.warehouse_name
                 == "mock_entity_two_wh"
             )