Skip to content

Commit

Permalink
[Tests] Fix an uncleaned DB file (#5589)
Browse files Browse the repository at this point in the history
  • Loading branch information
jond01 committed May 19, 2024
1 parent 5f8ae40 commit 1ca2503
Showing 1 changed file with 15 additions and 14 deletions.
29 changes: 15 additions & 14 deletions tests/model_monitoring/test_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,32 +11,32 @@
# 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.
#

import string
import time
import typing
import unittest.mock
from collections.abc import Iterator
from pathlib import Path
from random import choice, randint
from typing import Optional
from typing import Optional, cast

import pytest

import mlrun.common.schemas
import mlrun.model_monitoring
import mlrun.model_monitoring.db.stores.sqldb.sql_store
from mlrun.common.schemas.model_monitoring import ResultData, WriterEvent
from mlrun.model_monitoring.db.stores import ( # noqa: F401
StoreBase,
)
from mlrun.model_monitoring.db.stores.sqldb.sql_store import SQLStoreBase
from mlrun.model_monitoring.writer import _AppResultEvent

SQLStoreBase = typing.TypeVar("SQLStoreBase", bound="StoreBase")


class TestSQLStore:
_TEST_PROJECT = "test_model_endpoints"
_MODEL_ENDPOINT_ID = "some-ep-id"
_STORE_CONNECTION = "sqlite:///test.db"

@staticmethod
@pytest.fixture
def store_connection(tmp_path: Path) -> str:
return f"sqlite:///{tmp_path / 'test.db'}"

@pytest.fixture()
def _mock_random_endpoint(
Expand Down Expand Up @@ -107,15 +107,16 @@ def init_sql_tables(new_sql_store: SQLStoreBase):

@classmethod
@pytest.fixture
def new_sql_store(cls) -> SQLStoreBase:
def new_sql_store(cls, store_connection: str) -> Iterator[SQLStoreBase]:
# Generate store object target
store_type_object = mlrun.model_monitoring.db.ObjectStoreFactory(value="sql")
with unittest.mock.patch(
"mlrun.model_monitoring.helpers.get_connection_string",
return_value=cls._STORE_CONNECTION,
return_value=store_connection,
):
sql_store: SQLStoreBase = store_type_object.to_object_store(
project=cls._TEST_PROJECT
sql_store = cast(
SQLStoreBase,
store_type_object.to_object_store(project=cls._TEST_PROJECT),
)
yield sql_store
sql_store.delete_model_endpoints_resources()
Expand Down

0 comments on commit 1ca2503

Please sign in to comment.