Fixing apm

superagi/apm/analytics_helper.py

@@ -1,11 +1,10 @@
-from typing import List, Dict, Tuple, Optional, Iterator, Union, Any
+from typing import List, Dict, Union, Any
+
+from sqlalchemy import text, func, and_
 from sqlalchemy.orm import Session
-from sqlalchemy.orm.query import Query
+
 from superagi.models.events import Event
-from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy import text, func, Integer, and_
-from collections import defaultdict
-import logging
+
 
 class AnalyticsHelper:
 

superagi/apm/tools_handler.py

@@ -1,8 +1,11 @@
 from typing import List, Dict
-from sqlalchemy.orm import Session
-from sqlalchemy.orm.query import Query
+
 from sqlalchemy import func
+from sqlalchemy.orm import Session
+
 from superagi.models.events import Event
+
+
 class ToolsHandler:
 
     def __init__(self, session: Session, organisation_id: int):

tests/unit_tests/apm/test_analytics_helper.py

@@ -1,41 +1,38 @@
-import pytest
 from unittest.mock import MagicMock
+
+import pytest
+import superagi
 from superagi.apm.analytics_helper import AnalyticsHelper
-from sqlalchemy.orm import Session
 
-@pytest.fixture
-def mock_session():
-    return MagicMock(spec=Session)
 
 @pytest.fixture
 def organisation_id():
     return 1
 
+@pytest.fixture
+def mock_session():
+    return MagicMock()
+
 @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):
-    analytics_helper.calculate_run_completed_metrics = MagicMock(return_value = {})
+    mock_session.query().all.return_value = [MagicMock()]
     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):
-    analytics_helper.fetch_agent_data = MagicMock(return_value = {})
+    mock_session.query().all.return_value = [MagicMock()]
     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):
-    agent_id = 1
-    analytics_helper.fetch_agent_runs = MagicMock(return_value = [])
-    result = analytics_helper.fetch_agent_runs(agent_id)
+    mock_session.query().all.return_value = [MagicMock()]
+    result = analytics_helper.fetch_agent_runs(1)
     assert isinstance(result, list)
-    analytics_helper.fetch_agent_runs.assert_called_with(agent_id)
 
 def test_get_active_runs(analytics_helper, mock_session):
-    analytics_helper.get_active_runs = MagicMock(return_value = [])
+    mock_session.query().all.return_value = [MagicMock()]
     result = analytics_helper.get_active_runs()
-    assert isinstance(result, list)
-    analytics_helper.get_active_runs.assert_called()
+    assert isinstance(result, list)

tests/unit_tests/apm/test_tools_handler.py

@@ -1,22 +1,21 @@
 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(spec=Session)
 
 @pytest.fixture
 def organisation_id():
     return 1
 
+@pytest.fixture
+def mock_session():
+    return MagicMock()
+
 @pytest.fixture
 def tools_handler(mock_session, organisation_id):
     return ToolsHandler(mock_session, organisation_id)
 
-def test_calculate_tool_usage(tools_handler):
-    tools_handler.calculate_tool_usage = MagicMock(return_value=[])
+def test_calculate_tool_usage(tools_handler, mock_session):
+    mock_session.query().all.return_value = [MagicMock()]
     result = tools_handler.calculate_tool_usage()
-    assert isinstance(result, list)
-    tools_handler.calculate_tool_usage.assert_called()
+    assert isinstance(result, list)

tests/unit_tests/controllers/test_analytics.py

@@ -1,41 +1,63 @@
 from unittest.mock import patch
+from fastapi_jwt_auth import AuthJWT
+from superagi.helper.auth import get_user_organisation
 import pytest
 from fastapi.testclient import TestClient
 from main import app
 
 client = TestClient(app)
 
 def test_get_metrics_success():
-    with patch('superagi.controllers.analytics.AnalyticsHelper') as mock_helper:
+    with patch('superagi.helper.auth.get_user_organisation') as mock_org, \
+        patch('superagi.apm.analytics_helper.AnalyticsHelper') as mock_helper:
+
+        mock_org.return_value = 'dummy_organisation'
         mock_helper().calculate_run_completed_metrics.return_value = {'total_tokens': 10, 'total_calls': 5, 'runs_completed': 2}
+
         response = client.get("analytics/metrics")
         assert response.status_code == 200
         assert response.json() == {'total_tokens': 10, 'total_calls': 5, 'runs_completed': 2}
 
 def test_get_agents_success():
-    with patch('superagi.controllers.analytics.AnalyticsHelper') as mock_helper:
+    with patch('superagi.helper.auth.get_user_organisation') as mock_org, \
+        patch('superagi.apm.analytics_helper.AnalyticsHelper') as mock_helper:
+
+        mock_org.return_value = 'dummy_organisation'
         mock_helper().fetch_agent_data.return_value = {"agent_details": "mock_details", "model_info": "mock_info"}
+
         response = client.get("analytics/agents/all")
         assert response.status_code == 200
         assert response.json() == {"agent_details": "mock_details", "model_info": "mock_info"}
 
 def test_get_agent_runs_success():
-    with patch('superagi.controllers.analytics.AnalyticsHelper') as mock_helper:
+    with patch('superagi.helper.auth.get_user_organisation') as mock_org, \
+        patch('superagi.apm.analytics_helper.AnalyticsHelper') as mock_helper:
+
+        mock_org.return_value = 'dummy_organisation'
         mock_helper().fetch_agent_runs.return_value = "mock_agent_runs"
+
         response = client.get("analytics/agents/1")
         assert response.status_code == 200
         assert response.json() == "mock_agent_runs"
 
 def test_get_active_runs_success():
-    with patch('superagi.controllers.analytics.AnalyticsHelper') as mock_helper:
+    with patch('superagi.helper.auth.get_user_organisation') as mock_org, \
+        patch('superagi.apm.analytics_helper.AnalyticsHelper') as mock_helper:
+
+        mock_org.return_value = 'dummy_organisation'
         mock_helper().get_active_runs.return_value = ["mock_run_1", "mock_run_2"]
+
         response = client.get("analytics/runs/active")
         assert response.status_code == 200
         assert response.json() == ["mock_run_1", "mock_run_2"]
 
 def test_get_tools_user_success():
-    with patch('superagi.controllers.analytics.ToolsHandler') as mock_handler:
+    with patch('superagi.helper.auth.get_user_organisation') as mock_org, \
+        patch('superagi.apm.tools_handler.ToolsHandler') as mock_handler:
+
+        mock_org.return_value = 'dummy_organisation'
         mock_handler().calculate_tool_usage.return_value = ["tool1", "tool2"]
+
         response = client.get("analytics/tools/used")
         assert response.status_code == 200
         assert response.json() == ["tool1", "tool2"]