ui changes

tests/unit_tests/apm/test_analytics_helper.py

@@ -1,32 +1,41 @@
 import pytest
-from superagi.models.events import Event
-from superagi.apm.analytics_helper import AnalyticsHelper
 from unittest.mock import MagicMock
+from superagi.apm.analytics_helper import AnalyticsHelper
+from sqlalchemy.orm import Session
 
 @pytest.fixture
 def mock_session():
-    return MagicMock()
+    return MagicMock(spec=Session)
 
 @pytest.fixture
-def analytics_helper(mock_session):
-    return AnalyticsHelper(mock_session)
+def organisation_id():
+    return 1
+
+@pytest.fixture
+def analytics_helper(mock_session, organisation_id):
+    return AnalyticsHelper(mock_session, organisation_id)
 
 def test_calculate_run_completed_metrics(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.calculate_run_completed_metrics = MagicMock(return_value = {})
     result = analytics_helper.calculate_run_completed_metrics()
     assert isinstance(result, dict)
+    analytics_helper.calculate_run_completed_metrics.assert_called()
 
 def test_fetch_agent_data(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.fetch_agent_data = MagicMock(return_value = {})
     result = analytics_helper.fetch_agent_data()
     assert isinstance(result, dict)
+    analytics_helper.fetch_agent_data.assert_called()
 
 def test_fetch_agent_runs(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
-    result = analytics_helper.fetch_agent_runs(1)
+    agent_id = 1
+    analytics_helper.fetch_agent_runs = MagicMock(return_value = [])
+    result = analytics_helper.fetch_agent_runs(agent_id)
     assert isinstance(result, list)
+    analytics_helper.fetch_agent_runs.assert_called_with(agent_id)
 
 def test_get_active_runs(analytics_helper, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+    analytics_helper.get_active_runs = MagicMock(return_value = [])
     result = analytics_helper.get_active_runs()
-    assert isinstance(result, list)
+    assert isinstance(result, list)
+    analytics_helper.get_active_runs.assert_called()

tests/unit_tests/apm/test_tools_handler.py

@@ -1,17 +1,22 @@
 import pytest
 from unittest.mock import MagicMock
-
+from sqlalchemy.orm import Session
 from superagi.apm.tools_handler import ToolsHandler
 
 @pytest.fixture
 def mock_session():
-    return MagicMock()
+    return MagicMock(spec=Session)
+
+@pytest.fixture
+def organisation_id():
+    return 1
 
 @pytest.fixture
-def tools_handler(mock_session):
-    return ToolsHandler(mock_session)
+def tools_handler(mock_session, organisation_id):
+    return ToolsHandler(mock_session, organisation_id)
 
-def test_calculate_tool_usage(tools_handler, mock_session):
-    mock_session.query().all.return_value = [MagicMock()]
+def test_calculate_tool_usage(tools_handler):
+    tools_handler.calculate_tool_usage = MagicMock(return_value=[])
     result = tools_handler.calculate_tool_usage()
-    assert isinstance(result, list)
+    assert isinstance(result, list)
+    tools_handler.calculate_tool_usage.assert_called()