open-metadata--openmetadata/ingestion/tests/integration/ometa/test_ometa_mlmodel_api.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
OpenMetadata high-level API Model test
"""

import pytest

from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
from metadata.generated.schema.api.data.createDatabaseSchema import (
    CreateDatabaseSchemaRequest,
)
from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
from metadata.generated.schema.api.data.createTable import CreateTableRequest
from metadata.generated.schema.api.services.createDatabaseService import (
    CreateDatabaseServiceRequest,
)
from metadata.generated.schema.entity.data.mlmodel import (
    FeatureSource,
    FeatureSourceDataType,
    FeatureType,
    MlFeature,
    MlHyperParameter,
    MlModel,
)
from metadata.generated.schema.entity.data.table import Column, DataType, Table
from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
    BasicAuth,
)
from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
    MysqlConnection,
)
from metadata.generated.schema.entity.services.databaseService import (
    DatabaseConnection,
    DatabaseService,
    DatabaseServiceType,
)
from metadata.generated.schema.type.entityLineage import EntitiesEdge
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.generated.schema.type.entityReferenceList import EntityReferenceList

from .conftest import _safe_delete  # noqa: TID252


@pytest.fixture
def mlmodel_request(mlmodel_service):
    """Create ML model request using the ML model service from conftest."""
    return CreateMlModelRequest(
        name="test-model",
        algorithm="algo",
        service=mlmodel_service.fullyQualifiedName,
    )


@pytest.fixture
def expected_fqn(mlmodel_service):
    """Expected fully qualified name for test ML model."""
    return f"{mlmodel_service.name.root}.test-model"


class TestOMetaMlModelAPI:
    """
    ML Model API integration tests.
    Tests CRUD operations, versioning, ML features, hyperparameters, and lineage.

    Uses fixtures from conftest:
    - metadata: OpenMetadata client (session scope)
    - mlmodel_service: MlModelService (module scope)
    - create_user: User factory (function scope)
    - create_mlmodel: MlModel factory (function scope)
    """

    def test_create(self, metadata, mlmodel_service, mlmodel_request, expected_fqn, create_mlmodel):
        """
        We can create a Model and we receive it back as Entity
        """
        res = create_mlmodel(mlmodel_request)

        assert res.name.root == "test-model"
        assert res.algorithm == "algo"
        assert res.owners is None

        fetched = metadata.get_by_name(entity=MlModel, fqn=expected_fqn)
        assert fetched is not None
        assert fetched.id == res.id

    def test_update(
        self,
        metadata,
        mlmodel_service,
        mlmodel_request,
        expected_fqn,
        create_user,
        create_mlmodel,
    ):
        """
        Updating it properly changes its properties
        """
        user = create_user()
        owners = EntityReferenceList(root=[EntityReference(id=user.id, type="user")])

        res_create = create_mlmodel(mlmodel_request)

        updated = mlmodel_request.model_dump(exclude_unset=True)
        updated["owners"] = owners
        updated_entity = CreateMlModelRequest(**updated)

        res = metadata.create_or_update(data=updated_entity)

        assert res.algorithm == updated_entity.algorithm
        assert res_create.id == res.id
        assert res.owners.root[0].id == user.id

        res_none = metadata.get_by_name(entity=MlModel, fqn=expected_fqn)
        assert res_none.owners is None

        res_owner = metadata.get_by_name(
            entity=MlModel,
            fqn=expected_fqn,
            fields=["owners", "followers"],
        )
        assert res_owner.owners.root[0].id == user.id

    def test_get_name(self, metadata, mlmodel_request, expected_fqn, create_mlmodel):
        """
        We can fetch a model by name and get it back as Entity
        """
        created = create_mlmodel(mlmodel_request)

        res = metadata.get_by_name(entity=MlModel, fqn=expected_fqn)
        assert res.name.root == created.name.root

    def test_get_id(self, metadata, mlmodel_request, expected_fqn, create_mlmodel):
        """
        We can fetch a model by ID and get it back as Entity
        """
        create_mlmodel(mlmodel_request)

        res_name = metadata.get_by_name(entity=MlModel, fqn=expected_fqn)
        res = metadata.get_by_id(entity=MlModel, entity_id=res_name.id)

        assert res_name.id == res.id

    def test_list(self, metadata, mlmodel_request, create_mlmodel):
        """
        We can list all our models
        """
        created = create_mlmodel(mlmodel_request)

        res = metadata.list_entities(entity=MlModel)

        data = next(iter(ent for ent in res.entities if ent.name == created.name), None)
        assert data is not None

    def test_delete(self, metadata, mlmodel_request, expected_fqn, create_mlmodel):
        """
        We can delete a model by ID
        """
        created = create_mlmodel(mlmodel_request)

        metadata.delete(entity=MlModel, entity_id=str(created.id.root))

        deleted = metadata.get_by_name(entity=MlModel, fqn=expected_fqn)
        assert deleted is None

    def test_mlmodel_properties(self, metadata, mlmodel_service):
        """
        Check that we can create models with MLFeatures and MLHyperParams

        We can add lineage information
        """
        service = CreateDatabaseServiceRequest(
            name="test-service-table-ml",
            serviceType=DatabaseServiceType.Mysql,
            connection=DatabaseConnection(
                config=MysqlConnection(
                    username="username",
                    authType=BasicAuth(
                        password="password",
                    ),
                    hostPort="http://localhost:1234",
                )
            ),
        )
        service_entity = metadata.create_or_update(data=service)

        create_db = CreateDatabaseRequest(
            name="test-db-ml",
            service=service_entity.fullyQualifiedName,
        )
        create_db_entity = metadata.create_or_update(data=create_db)

        create_schema = CreateDatabaseSchemaRequest(
            name="test-schema-ml",
            database=create_db_entity.fullyQualifiedName,
        )
        create_schema_entity = metadata.create_or_update(data=create_schema)

        create_table1 = CreateTableRequest(
            name="test-ml",
            databaseSchema=create_schema_entity.fullyQualifiedName,
            columns=[Column(name="education", dataType=DataType.STRING)],
        )
        table1_entity = metadata.create_or_update(data=create_table1)

        create_table2 = CreateTableRequest(
            name="another_test-ml",
            databaseSchema=create_schema_entity.fullyQualifiedName,
            columns=[Column(name="age", dataType=DataType.INT)],
        )
        table2_entity = metadata.create_or_update(data=create_table2)

        model = CreateMlModelRequest(
            name="test-model-lineage",
            algorithm="algo",
            mlFeatures=[
                MlFeature(
                    name="age",
                    dataType=FeatureType.numerical,
                    featureSources=[
                        FeatureSource(
                            name="age",
                            dataType=FeatureSourceDataType.integer,
                            dataSource=metadata.get_entity_reference(
                                entity=Table, fqn=table2_entity.fullyQualifiedName
                            ),
                        )
                    ],
                ),
                MlFeature(
                    name="persona",
                    dataType=FeatureType.categorical,
                    featureSources=[
                        FeatureSource(
                            name="age",
                            dataType=FeatureSourceDataType.integer,
                            dataSource=metadata.get_entity_reference(
                                entity=Table, fqn=table2_entity.fullyQualifiedName
                            ),
                        ),
                        FeatureSource(
                            name="education",
                            dataType=FeatureSourceDataType.string,
                            dataSource=metadata.get_entity_reference(
                                entity=Table, fqn=table1_entity.fullyQualifiedName
                            ),
                        ),
                        FeatureSource(name="city", dataType=FeatureSourceDataType.string),
                    ],
                    featureAlgorithm="PCA",
                ),
            ],
            mlHyperParameters=[
                MlHyperParameter(name="regularisation", value="0.5"),
                MlHyperParameter(name="random", value="hello"),
            ],
            target="myTarget",
            service=mlmodel_service.fullyQualifiedName,
        )

        res = metadata.create_or_update(data=model)

        try:
            assert res.mlFeatures is not None
            assert res.mlHyperParameters is not None

            lineage = metadata.get_lineage_by_id(entity=MlModel, entity_id=str(res.id.root))

            nodes = {node["id"] for node in lineage["nodes"]}
            assert nodes == {str(table1_entity.id.root), str(table2_entity.id.root)}

            for edge in lineage.get("upstreamEdges") or []:
                metadata.delete_lineage_edge(
                    edge=EntitiesEdge(
                        fromEntity=EntityReference(id=edge["fromEntity"], type="table"),
                        toEntity=EntityReference(id=edge["toEntity"], type="mlmodel"),
                    )
                )

            metadata.add_mlmodel_lineage(model=res)

            lineage = metadata.get_lineage_by_id(entity=MlModel, entity_id=str(res.id.root))

            nodes = {node["id"] for node in lineage["nodes"]}
            assert nodes == {str(table1_entity.id.root), str(table2_entity.id.root)}
        finally:
            _safe_delete(metadata, entity=MlModel, entity_id=res.id, hard_delete=True)
            _safe_delete(
                metadata,
                entity=DatabaseService,
                entity_id=service_entity.id,
                recursive=True,
                hard_delete=True,
            )

    def test_list_versions(self, metadata, mlmodel_request, create_mlmodel):
        """
        Test listing ML model entity versions
        """
        created = create_mlmodel(mlmodel_request)

        res = metadata.get_list_entity_versions(entity=MlModel, entity_id=created.id.root)
        assert res is not None
        assert len(res.versions) >= 1

    def test_get_entity_version(self, metadata, mlmodel_request, create_mlmodel):
        """
        Test retrieving a specific ML model entity version
        """
        created = create_mlmodel(mlmodel_request)

        res = metadata.get_entity_version(entity=MlModel, entity_id=created.id.root, version=0.1)

        assert res.version.root == 0.1
        assert res.id == created.id

    def test_get_entity_ref(self, metadata, mlmodel_request, create_mlmodel):
        """
        Test retrieving EntityReference for an ML model
        """
        created = create_mlmodel(mlmodel_request)
        entity_ref = metadata.get_entity_reference(entity=MlModel, fqn=created.fullyQualifiedName)

        assert created.id == entity_ref.id