refactor(api): extract workflow statistic query service (#41271)

api/tests/unit_tests/controllers/console/app/test_app_apis.py

@@ -5,13 +5,14 @@ from __future__ import annotations
 import uuid
 from collections.abc import Iterator
 from inspect import unwrap
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 from flask import Flask
 from pydantic import ValidationError
 from sqlalchemy import Engine
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, NotFound
 
 from controllers.console import console_ns
@@ -70,7 +71,6 @@ from models import App, Site
 from models.account import Account, AccountStatus
 from models.engine import db
 from models.trigger import WorkflowWebhookTrigger
-from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository
 from services.app_site_service import (
     AppSiteAppNotFoundError,
     AppSiteChanges,
@@ -106,6 +106,15 @@ def _make_app(
     return app
 
 
+def _make_request_context() -> RequestContext:
+    return RequestContext(
+        request_id="request-1",
+        trace_id="trace-1",
+        account_id=USER_ID,
+        active_workspace_id=TENANT_ID,
+    )
+
+
 @pytest.fixture
 def database_app() -> Iterator[Flask]:
     app = Flask(__name__)
@@ -690,48 +699,162 @@ class TestWorkflowStatisticEndpoints:
         assert query.end is None
 
     def test_workflow_daily_runs_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
-        repository = DifyAPISQLAlchemyWorkflowRunRepository(
-            session_maker=sessionmaker(bind=db.engine, expire_on_commit=False)
-        )
-        monkeypatch.setattr(
-            repository,
-            "get_daily_runs_statistics",
-            lambda **_kwargs: [{"date": "2024-01-01"}],
-        )
-        monkeypatch.setattr(
-            workflow_statistic_module.DifyAPIRepositoryFactory,
-            "create_api_workflow_run_repository",
-            lambda *_args, **_kwargs: repository,
-        )
+        service = MagicMock()
+        service.get_daily_runs.return_value = [{"date": "2024-01-01", "runs": 2}]
         monkeypatch.setattr(
             workflow_statistic_module,
-            "parse_time_range",
-            lambda *_args, **_kwargs: (None, None),
+            "application_services",
+            lambda: SimpleNamespace(workflow_statistics=service),
         )
+        self._patch_statistic_time_range(monkeypatch)
 
         api = workflow_statistic_module.WorkflowDailyRunsStatistic()
         method = unwrap(api.get)
+        request_context = _make_request_context()
 
         with database_app.test_request_context("/"):
-            account = _make_account()
-            account.timezone = "UTC"
-            response = method(api, WorkflowStatisticQuery(), account, app_model=_make_app("app-1", tenant_id="t1"))
+            response = method(
+                api,
+                WorkflowStatisticQuery(),
+                request_context,
+                app_model=_make_app("app-1", tenant_id="t1"),
+            )
 
-        assert response.get_json() == {"data": [{"date": "2024-01-01"}]}
+        assert response == {"data": [{"date": "2024-01-01", "runs": 2}]}
+        service.get_daily_runs.assert_called_once_with(
+            request_context,
+            app_id="app-1",
+            start_date=None,
+            end_date=None,
+            timezone="UTC",
+        )
 
     def test_workflow_daily_terminals_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
-        repository = DifyAPISQLAlchemyWorkflowRunRepository(
-            session_maker=sessionmaker(bind=db.engine, expire_on_commit=False)
+        service = MagicMock()
+        service.get_daily_terminals.return_value = [{"date": "2024-01-02", "terminal_count": 3}]
+        monkeypatch.setattr(
+            workflow_statistic_module,
+            "application_services",
+            lambda: SimpleNamespace(workflow_statistics=service),
+        )
+        self._patch_statistic_time_range(monkeypatch)
+
+        api = workflow_statistic_module.WorkflowDailyTerminalsStatistic()
+        method = unwrap(api.get)
+        request_context = _make_request_context()
+
+        with database_app.test_request_context("/"):
+            response = method(
+                api,
+                WorkflowStatisticQuery(),
+                request_context,
+                app_model=_make_app("app-1", tenant_id="t1"),
+            )
+
+        assert response == {"data": [{"date": "2024-01-02", "terminal_count": 3}]}
+        service.get_daily_terminals.assert_called_once_with(
+            request_context,
+            app_id="app-1",
+            start_date=None,
+            end_date=None,
+            timezone="UTC",
+        )
+
+    def test_workflow_daily_token_cost_statistic(
+        self,
+        database_app: Flask,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        service = MagicMock()
+        service.get_daily_token_costs.return_value = [{"date": "2024-01-03", "token_count": 4}]
+        monkeypatch.setattr(
+            workflow_statistic_module,
+            "application_services",
+            lambda: SimpleNamespace(workflow_statistics=service),
+        )
+        self._patch_statistic_time_range(monkeypatch)
+
+        api = workflow_statistic_module.WorkflowDailyTokenCostStatistic()
+        method = unwrap(api.get)
+        request_context = _make_request_context()
+
+        with database_app.test_request_context("/"):
+            response = method(
+                api,
+                WorkflowStatisticQuery(),
+                request_context,
+                app_model=_make_app("app-1", tenant_id="t1"),
+            )
+
+        assert response == {"data": [{"date": "2024-01-03", "token_count": 4}]}
+        service.get_daily_token_costs.assert_called_once_with(
+            request_context,
+            app_id="app-1",
+            start_date=None,
+            end_date=None,
+            timezone="UTC",
+        )
+
+    def test_workflow_average_app_interaction_statistic(
+        self,
+        database_app: Flask,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        service = MagicMock()
+        service.get_average_app_interactions.return_value = [{"date": "2024-01-04", "interactions": 2.5}]
+        monkeypatch.setattr(
+            workflow_statistic_module,
+            "application_services",
+            lambda: SimpleNamespace(workflow_statistics=service),
+        )
+        self._patch_statistic_time_range(monkeypatch)
+
+        api = workflow_statistic_module.WorkflowAverageAppInteractionStatistic()
+        method = unwrap(api.get)
+        request_context = _make_request_context()
+
+        with database_app.test_request_context("/"):
+            response = method(
+                api,
+                WorkflowStatisticQuery(),
+                request_context,
+                app_model=_make_app("app-1", tenant_id="t1"),
+            )
+
+        assert response == {"data": [{"date": "2024-01-04", "interactions": 2.5}]}
+        service.get_average_app_interactions.assert_called_once_with(
+            request_context,
+            app_id="app-1",
+            start_date=None,
+            end_date=None,
+            timezone="UTC",
+        )
+
+    def test_workflow_statistic_invalid_time_range(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        account = _make_account()
+        account.timezone = "UTC"
+        monkeypatch.setattr(
+            workflow_statistic_module,
+            "current_account_with_tenant",
+            lambda: SimpleNamespace(account=account),
         )
         monkeypatch.setattr(
-            repository,
-            "get_daily_terminals_statistics",
-            lambda **_kwargs: [{"date": "2024-01-02"}],
+            workflow_statistic_module,
+            "parse_time_range",
+            MagicMock(side_effect=ValueError("invalid range")),
         )
+
+        with pytest.raises(BadRequest, match="invalid range"):
+            workflow_statistic_module._resolve_statistic_time_range(WorkflowStatisticQuery())
+
+    @staticmethod
+    def _patch_statistic_time_range(monkeypatch: pytest.MonkeyPatch) -> None:
+        account = _make_account()
+        account.timezone = "UTC"
         monkeypatch.setattr(
-            workflow_statistic_module.DifyAPIRepositoryFactory,
-            "create_api_workflow_run_repository",
-            lambda *_args, **_kwargs: repository,
+            workflow_statistic_module,
+            "current_account_with_tenant",
+            lambda: SimpleNamespace(account=account),
         )
         monkeypatch.setattr(
             workflow_statistic_module,
@@ -739,16 +862,6 @@ class TestWorkflowStatisticEndpoints:
             lambda *_args, **_kwargs: (None, None),
         )
 
-        api = workflow_statistic_module.WorkflowDailyTerminalsStatistic()
-        method = unwrap(api.get)
-
-        with database_app.test_request_context("/"):
-            account = _make_account()
-            account.timezone = "UTC"
-            response = method(api, WorkflowStatisticQuery(), account, app_model=_make_app("app-1", tenant_id="t1"))
-
-        assert response.get_json() == {"data": [{"date": "2024-01-02"}]}
-
 
 class TestWorkflowTriggerEndpoints:
     def test_webhook_trigger_payload(self):

api/tests/unit_tests/extensions/test_ext_application_services.py

@@ -41,6 +41,7 @@ from services.init_validation_service import InvalidInitializationPasswordError
 from services.partner_tenant_binding_service import PartnerTenantBindingService
 from services.tag_application_service import TagApplicationService
 from services.webapp_access_query_service import WebAppAccessUnavailableError
+from services.workflow_statistic_query_service import WorkflowStatisticQueryService
 
 
 @pytest.mark.parametrize(
@@ -111,6 +112,7 @@ def test_init_app_registers_services_for_the_current_app(
         services = ext_application_services.application_services()
         assert services is app.extensions["application_services"]
         assert services.init_validation.is_validated(session_validated=False) is False
+        assert isinstance(services.workflow_statistics, WorkflowStatisticQueryService)
 
 
 @pytest.mark.parametrize(

api/tests/unit_tests/services/test_workflow_statistic_query_service.py

@@ -0,0 +1,88 @@
+from datetime import UTC, datetime
+from unittest.mock import MagicMock
+
+import pytest
+
+from machinery.context import RequestContext
+from models.enums import WorkflowRunTriggeredFrom
+from repositories.api_workflow_run_repository import APIWorkflowRunRepository
+from services.workflow_statistic_query_service import WorkflowStatisticQueryService
+
+
+def _request_context(*, workspace_id: str | None = "workspace-1") -> RequestContext:
+    return RequestContext(
+        request_id="request-1",
+        trace_id="trace-1",
+        account_id="account-1",
+        active_workspace_id=workspace_id,
+    )
+
+
+def test_workflow_statistic_queries_delegate_to_workflow_run_repository() -> None:
+    workflow_runs = MagicMock(spec=APIWorkflowRunRepository)
+    workflow_runs.get_daily_runs_statistics.return_value = [{"date": "2024-01-01", "runs": 2}]
+    workflow_runs.get_daily_terminals_statistics.return_value = [{"date": "2024-01-01", "terminal_count": 3}]
+    workflow_runs.get_daily_token_cost_statistics.return_value = [{"date": "2024-01-01", "token_count": 4}]
+    workflow_runs.get_average_app_interaction_statistics.return_value = [{"date": "2024-01-01", "interactions": 2.5}]
+    service = WorkflowStatisticQueryService(workflow_runs=workflow_runs)
+    context = _request_context()
+    start_date = datetime(2024, 1, 1, tzinfo=UTC)
+    end_date = datetime(2024, 1, 2, tzinfo=UTC)
+
+    assert service.get_daily_runs(
+        context,
+        app_id="app-1",
+        start_date=start_date,
+        end_date=end_date,
+        timezone="Asia/Shanghai",
+    ) == [{"date": "2024-01-01", "runs": 2}]
+    assert service.get_daily_terminals(
+        context,
+        app_id="app-1",
+        start_date=start_date,
+        end_date=end_date,
+        timezone="Asia/Shanghai",
+    ) == [{"date": "2024-01-01", "terminal_count": 3}]
+    assert service.get_daily_token_costs(
+        context,
+        app_id="app-1",
+        start_date=start_date,
+        end_date=end_date,
+        timezone="Asia/Shanghai",
+    ) == [{"date": "2024-01-01", "token_count": 4}]
+    assert service.get_average_app_interactions(
+        context,
+        app_id="app-1",
+        start_date=start_date,
+        end_date=end_date,
+        timezone="Asia/Shanghai",
+    ) == [{"date": "2024-01-01", "interactions": 2.5}]
+
+    expected_arguments = {
+        "tenant_id": "workspace-1",
+        "app_id": "app-1",
+        "triggered_from": WorkflowRunTriggeredFrom.APP_RUN,
+        "start_date": start_date,
+        "end_date": end_date,
+        "timezone": "Asia/Shanghai",
+    }
+    workflow_runs.get_daily_runs_statistics.assert_called_once_with(**expected_arguments)
+    workflow_runs.get_daily_terminals_statistics.assert_called_once_with(**expected_arguments)
+    workflow_runs.get_daily_token_cost_statistics.assert_called_once_with(**expected_arguments)
+    workflow_runs.get_average_app_interaction_statistics.assert_called_once_with(**expected_arguments)
+
+
+def test_workflow_statistic_query_requires_active_workspace() -> None:
+    workflow_runs = MagicMock(spec=APIWorkflowRunRepository)
+    service = WorkflowStatisticQueryService(workflow_runs=workflow_runs)
+
+    with pytest.raises(RuntimeError, match="did not resolve an active workspace"):
+        service.get_daily_runs(
+            _request_context(workspace_id=None),
+            app_id="app-1",
+            start_date=None,
+            end_date=None,
+            timezone="UTC",
+        )
+
+    workflow_runs.get_daily_runs_statistics.assert_not_called()