1  import json
  2  import os
  3  from typing import Optional
  4  
  5  import pytest
  6  from litestar import Litestar
  7  from litestar import get
  8  from litestar.datastructures import State
  9  from litestar.testing import TestClient
 10  
 11  from evidently._pydantic_compat import BaseModel
 12  from evidently.legacy.ui.app import create_app
 13  from evidently.legacy.ui.base import Project
 14  from evidently.legacy.ui.components.base import Component
 15  from evidently.legacy.ui.components.base import ComponentContext
 16  from evidently.legacy.ui.components.storage import LocalStorageComponent
 17  from evidently.legacy.ui.local_service import LocalConfig
 18  from evidently.legacy.ui.local_service import LocalServiceComponent
 19  from evidently.legacy.ui.managers.projects import ProjectManager
 20  from evidently.legacy.ui.security.service import SecurityService
 21  from evidently.legacy.utils import NumpyEncoder
 22  
 23  HEADERS = {"Content-Type": "application/json"}
 24  
 25  os.environ["DO_NOT_TRACK"] = "1"
 26  
 27  
 28  class TestsSetupComponent(Component):
 29      class Config:
 30          arbitrary_types_allowed = True
 31  
 32      app: Optional[Litestar] = None
 33  
 34      def get_route_handlers(self, ctx: ComponentContext):
 35          @get("/tests_setup")
 36          async def tests_setup(project_manager: ProjectManager, security: SecurityService) -> None:
 37              self.app.state["pm"] = project_manager
 38              self.app.state["security"] = security
 39  
 40          return [tests_setup]
 41  
 42      def finalize(self, ctx: ComponentContext, app: Litestar):
 43          self.app = app
 44          client = TestClient(app)
 45          client.get("/tests_setup")
 46  
 47  
 48  def _get_app(app) -> Litestar:
 49      if isinstance(app, Litestar):
 50          return app
 51      return _get_app(app.__closure__[0].cell_contents.app)
 52  
 53  
 54  def _get_app_state(client: TestClient) -> State:
 55      try:
 56          return _get_app(client.app).state
 57      except (AttributeError, IndexError) as e:
 58          raise ValueError("Cannot find app state") from e
 59  
 60  
 61  def get_pm(client: TestClient) -> ProjectManager:
 62      try:
 63          return _get_app_state(client)["pm"]
 64      except KeyError as e:
 65          raise ValueError("Cannot find pm in client") from e
 66  
 67  
 68  def get_security(client: TestClient) -> SecurityService:
 69      try:
 70          return _get_app_state(client)["security"]
 71      except KeyError as e:
 72          raise ValueError("Cannot find security in client") from e
 73  
 74  
 75  def _dumps(obj: BaseModel):
 76      return json.dumps(obj.dict(), allow_nan=True, cls=NumpyEncoder)
 77  
 78  
 79  @pytest.fixture
 80  def test_client(tmp_path):
 81      config = LocalConfig(
 82          storage=LocalStorageComponent(path=str(tmp_path)),
 83          service=LocalServiceComponent(debug=True),
 84          additional_components={"_setup_tests": TestsSetupComponent()},
 85      )
 86      return TestClient(create_app(config=config))
 87  
 88  
 89  @pytest.fixture
 90  def project_manager(test_client) -> ProjectManager:
 91      return get_pm(test_client)
 92  
 93  
 94  @pytest.fixture
 95  def mock_project():
 96      return Project(name="mock", team_id=None)
 97  
 98  
 99  @pytest.fixture
100  def project_factory():
101      def inner(name: str) -> Project:
102          return Project(name=name, team_id=None)
103  
104      return inner