mito-ai 0.1.50__py3-none-any.whl

mito_ai/tests/rules/rules_test.py

@@ -0,0 +1,117 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+
+RULE_NAME = "test_rule"
+RULE_CONTENT = "This is a test rule for data analysis."
+
+from mito_ai.tests.conftest import TOKEN
+
+# --- PUT RULES ---
+
+
+def test_put_rule_with_auth(jp_base_url):
+    response = requests.put(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        headers={"Authorization": f"token {TOKEN}"},
+        json={"content": RULE_CONTENT},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["status"] == "updated"
+    assert response_json["rules file "] == RULE_NAME
+
+
+def test_put_rule_with_no_auth(jp_base_url):
+    response = requests.put(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        json={"content": RULE_CONTENT},
+    )
+
+    assert response.status_code == 403  # Forbidden
+
+
+def test_put_rule_with_incorrect_auth(jp_base_url):
+    response = requests.put(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        headers={"Authorization": f"token incorrect-token"}, # <- wrong token
+        json={"content": RULE_CONTENT},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_put_rule_missing_content(jp_base_url):
+    response = requests.put(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        headers={"Authorization": f"token {TOKEN}"},
+        json={},  # Missing content field
+    )
+    assert response.status_code == 400  # Bad Request
+
+
+# --- GET SPECIFIC RULE ---
+
+
+def test_get_rule_with_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    assert response.json() == {"key": RULE_NAME, "content": RULE_CONTENT}
+
+
+def test_get_rule_with_no_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_rule_with_incorrect_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules/{RULE_NAME}",
+        headers={"Authorization": f"token incorrect-token"}, # <- wrong token
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_nonexistent_rule_with_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules/nonexistent_rule",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 404  # Not Found
+
+
+# --- GET ALL RULES ---
+
+
+def test_get_all_rules_with_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    
+    response_json = response.json()
+    assert isinstance(response_json, list)
+    # Should contain our test rule (with .md extension)
+    assert f"{RULE_NAME}.md" in response_json
+
+
+def test_get_all_rules_with_no_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_all_rules_with_incorrect_auth(jp_base_url):
+    response = requests.get(
+        jp_base_url + f"/mito-ai/rules",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/server_limits_test.py

@@ -0,0 +1,406 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from datetime import datetime
+from typing import Dict, Optional
+import pytest
+from unittest.mock import patch, MagicMock, call
+
+from mito_ai.completions.models import MessageType
+from mito_ai.utils.server_limits import (
+    OS_MONTHLY_AI_COMPLETIONS_LIMIT,
+    OS_MONTHLY_AUTOCOMPLETE_LIMIT,
+    check_mito_server_quota,
+    update_mito_server_quota,
+    UJ_AI_MITO_API_NUM_USAGES,
+    UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES,
+    UJ_MITO_AI_LAST_RESET_DATE,
+    UJ_MITO_AI_FIRST_USAGE_DATE
+)
+from mito_ai.utils.telemetry_utils import MITO_SERVER_FREE_TIER_LIMIT_REACHED
+
+# Constants for testing
+CURRENT_DATE = datetime.now().strftime("%Y-%m-%d")
+OLD_DATE = "2020-01-01"
+FUTURE_DATE = (datetime.now().replace(year=datetime.now().year + 1)).strftime("%Y-%m-%d")
+
+# Each test case is a tuple of:
+# (is_pro, chat_completion_count, autocomplete_count, last_reset_date, message_type, should_raise_error)
+
+# 1. Open Source Subscription Tests
+OS_SUBSCRIPTION_TESTS = [
+    # Below both limits
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.CHAT, False),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+    
+    # Above AI completions limit only
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.CHAT, True),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+    
+    # Above autocomplete limit only
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1, CURRENT_DATE, MessageType.CHAT, False),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, True),
+    
+    # Above both limits
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1, CURRENT_DATE, MessageType.CHAT, True),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, True),
+    
+    # Equal to limits
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.CHAT, True),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT, CURRENT_DATE, MessageType.CHAT, False),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT, CURRENT_DATE, MessageType.INLINE_COMPLETION, True),
+]
+
+# 2. Pro Subscription Tests
+PRO_SUBSCRIPTION_TESTS = [
+    # Pro user below limits
+    (True, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.CHAT, False),
+    (True, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+    
+    # Pro user above limits - should still not raise error
+    (True, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 100, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 100, CURRENT_DATE, MessageType.CHAT, False),
+    (True, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 100, OS_MONTHLY_AUTOCOMPLETE_LIMIT + 100, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+]
+
+# 3. Error Handling Tests - Note: None values are replaced in the test function
+ERROR_HANDLING_TESTS = [
+    # Missing last reset date - should not raise error and set the date
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, None, MessageType.CHAT, False),
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT - 1, OS_MONTHLY_AUTOCOMPLETE_LIMIT - 1, None, MessageType.INLINE_COMPLETION, False),
+]
+
+# 4. Edge Cases
+EDGE_CASE_TESTS = [
+    # Negative counts - should be treated as below limit
+    (False, -10, -10, CURRENT_DATE, MessageType.CHAT, False),
+    (False, -10, -10, CURRENT_DATE, MessageType.INLINE_COMPLETION, False),
+]
+
+# Combine all tests
+ALL_TEST_CASES = OS_SUBSCRIPTION_TESTS + PRO_SUBSCRIPTION_TESTS + ERROR_HANDLING_TESTS + EDGE_CASE_TESTS
+
+@pytest.mark.parametrize(
+    "is_pro, chat_completion_count, autocomplete_count, last_reset_date, message_type, should_raise_error",
+    ALL_TEST_CASES
+)
+def test_check_mito_server_quota(
+    is_pro: bool, 
+    chat_completion_count: int,
+    autocomplete_count: int,
+    last_reset_date: Optional[str], 
+    message_type: MessageType, 
+    should_raise_error: bool
+) -> None:
+    """Test the check_mito_server_quota function with various combinations of inputs."""
+    
+    # Create the patch context managers
+    patches = [
+        patch("mito_ai.utils.server_limits.is_pro", return_value=is_pro),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=chat_completion_count),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=autocomplete_count),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=last_reset_date),
+    ]
+    
+    # Add set_user_field mock if needed for None values
+    if last_reset_date is None:
+        patches.append(patch("mito_ai.utils.server_limits.set_user_field"))
+    
+    # Apply all patches
+    for p in patches:
+        p.start()
+    
+    try:
+        if should_raise_error:
+            # Should raise PermissionError
+            with pytest.raises(PermissionError) as exc_info:
+                check_mito_server_quota(message_type)
+            assert str(exc_info.value) == MITO_SERVER_FREE_TIER_LIMIT_REACHED
+        else:
+            # Should not raise any error
+            check_mito_server_quota(message_type)
+    finally:
+        # Stop all patches
+        for p in patches:
+            p.stop()
+
+# New dedicated test cases for date-related reset behavior
+# Each test case is a tuple of:
+# (reset_date, is_pro, message_type, should_reset)
+DATE_RESET_TEST_CASES = [
+    # Old date should trigger reset for both message types
+    (OLD_DATE, False, MessageType.CHAT, True),
+    (OLD_DATE, False, MessageType.INLINE_COMPLETION, True),
+    
+    # Pro user with old date should not trigger reset (they have no limits)
+    (OLD_DATE, True, MessageType.CHAT, False),
+    (OLD_DATE, True, MessageType.INLINE_COMPLETION, False),
+    
+    # Current date should not trigger reset
+    (CURRENT_DATE, False, MessageType.CHAT, False),
+    (CURRENT_DATE, False, MessageType.INLINE_COMPLETION, False),
+    
+    # Future date should not trigger reset
+    (FUTURE_DATE, False, MessageType.CHAT, False),
+    (FUTURE_DATE, False, MessageType.INLINE_COMPLETION, False),
+]
+
+@pytest.mark.parametrize(
+    "reset_date, is_pro, message_type, should_reset",
+    DATE_RESET_TEST_CASES
+)
+def test_date_triggers_reset(reset_date: Optional[str], is_pro: bool, message_type: MessageType, should_reset: bool) -> None:
+    """
+    Test whether different dates trigger counter reset operations.
+    Rather than checking downstream behavior, directly verify the reset operation occurs.
+    """
+    # Mock set_user_field to track calls
+    mock_set_user_field = MagicMock()
+    
+    with (
+        patch("mito_ai.utils.server_limits.is_pro", return_value=is_pro),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=reset_date),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=1),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=1),
+        patch("mito_ai.utils.server_limits.set_user_field", mock_set_user_field)
+    ):
+        # Call the function
+        check_mito_server_quota(message_type)
+        
+        # Verify whether reset was triggered based on expectations
+        reset_calls = [
+            call(UJ_AI_MITO_API_NUM_USAGES, 0),
+            call(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+        ]
+        
+        if should_reset:
+            # Check that the API usages and autocomplete usages were reset to 0
+            # The format is slightly different because of how we're building the calls list
+            for reset_call in reset_calls:
+                assert reset_call in mock_set_user_field.call_args_list, f"Expected reset call {reset_call} not found"
+            
+            # Also verify that the reset date was updated
+            assert any(call[0][0] == UJ_MITO_AI_LAST_RESET_DATE for call in mock_set_user_field.call_args_list), "Reset date was not updated"
+        else:
+            # For cases that shouldn't reset, verify the reset calls were NOT made
+            # If we're setting a missing date (None), we'll still see calls, but they won't be resets
+            if reset_date is not None:
+                for reset_call in reset_calls:
+                    assert reset_call not in mock_set_user_field.call_args_list, f"Unexpected reset call {reset_call} found"
+
+# Special test cases that require specific patching of constants
+def test_zero_limits_configured() -> None:
+    """Test when limits are configured as zero."""
+    with (
+        patch("mito_ai.utils.server_limits.OS_MONTHLY_AI_COMPLETIONS_LIMIT", 0),
+        patch("mito_ai.utils.server_limits.OS_MONTHLY_AUTOCOMPLETE_LIMIT", 0),
+        patch("mito_ai.utils.server_limits.is_pro", return_value=False),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=1),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=1),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=CURRENT_DATE)
+    ):
+        # Test with chat completion - should raise error
+        with pytest.raises(PermissionError) as exc_info:
+            check_mito_server_quota(MessageType.CHAT)
+        assert str(exc_info.value) == MITO_SERVER_FREE_TIER_LIMIT_REACHED
+        
+        # Test with inline completion - should raise error
+        with pytest.raises(PermissionError) as exc_info:
+            check_mito_server_quota(MessageType.INLINE_COMPLETION)
+        assert str(exc_info.value) == MITO_SERVER_FREE_TIER_LIMIT_REACHED
+
+# New test cases for update_mito_server_quota function
+# Each case is (first_usage_date, last_reset_date, current_completion_count, current_autocomplete_count, message_type, expected_operations)
+# The expected_operations is a dict specifying what operations should happen: reset_counters, update_first_date, increment_chat, increment_autocomplete
+UPDATE_QUOTA_TEST_CASES = [
+    # Basic cases - incrementing counters
+    (
+        CURRENT_DATE, CURRENT_DATE, 10, 20, MessageType.CHAT, 
+        {"reset_counters": False, "update_first_date": False, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    (
+        CURRENT_DATE, CURRENT_DATE, 10, 20, MessageType.INLINE_COMPLETION, 
+        {"reset_counters": False, "update_first_date": False, "increment_chat": False, "increment_autocomplete": True}
+    ),
+    
+    # First usage initialization
+    (
+        None, CURRENT_DATE, 10, 20, MessageType.CHAT, 
+        {"reset_counters": False, "update_first_date": True, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    
+    # Reset date initialization
+    (
+        CURRENT_DATE, None, 10, 20, MessageType.CHAT, 
+        {"reset_counters": False, "update_first_date": False, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    
+    # Both dates missing
+    (
+        None, None, 10, 20, MessageType.CHAT, 
+        {"reset_counters": False, "update_first_date": True, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    
+    # Old reset date - should reset counters before updating
+    (
+        CURRENT_DATE, OLD_DATE, 10, 20, MessageType.CHAT, 
+        {"reset_counters": True, "update_first_date": False, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    (
+        CURRENT_DATE, OLD_DATE, 10, 20, MessageType.INLINE_COMPLETION, 
+        {"reset_counters": True, "update_first_date": False, "increment_chat": False, "increment_autocomplete": True}
+    ),
+    
+    # None counts - should be handled as 0
+    (
+        CURRENT_DATE, CURRENT_DATE, None, 20, MessageType.CHAT, 
+        {"reset_counters": False, "update_first_date": False, "increment_chat": True, "increment_autocomplete": False}
+    ),
+    (
+        CURRENT_DATE, CURRENT_DATE, 10, None, MessageType.INLINE_COMPLETION, 
+        {"reset_counters": False, "update_first_date": False, "increment_chat": False, "increment_autocomplete": True}
+    ),
+]
+
+@pytest.mark.parametrize(
+    "first_usage_date, last_reset_date, completion_count, autocomplete_count, message_type, expected_ops",
+    UPDATE_QUOTA_TEST_CASES
+)
+def test_update_mito_server_quota(
+    first_usage_date: Optional[str], 
+    last_reset_date: Optional[str], 
+    completion_count: int, 
+    autocomplete_count: int, 
+    message_type: MessageType, 
+    expected_ops: Dict[str, bool]
+) -> None:
+    """Test the update_mito_server_quota function with various input combinations."""
+    # Mock set_user_field to track calls
+    mock_set_user_field = MagicMock()
+    
+    with (
+        patch("mito_ai.utils.server_limits.get_first_completion_date", return_value=first_usage_date),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=last_reset_date),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=completion_count),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=autocomplete_count),
+        patch("mito_ai.utils.server_limits.set_user_field", mock_set_user_field)
+    ):
+        # Call the function
+        update_mito_server_quota(message_type)
+        
+        # Check initialization of first usage date
+        if first_usage_date is None or expected_ops["update_first_date"]:
+            assert any(call[0][0] == UJ_MITO_AI_FIRST_USAGE_DATE for call in mock_set_user_field.call_args_list), \
+                "First usage date should have been initialized"
+        
+        # Check initialization of reset date
+        if last_reset_date is None:
+            assert any(call[0][0] == UJ_MITO_AI_LAST_RESET_DATE for call in mock_set_user_field.call_args_list), \
+                "Reset date should have been initialized"
+        
+        # Check counter reset operations
+        if expected_ops["reset_counters"]:
+            reset_calls = [
+                call(UJ_AI_MITO_API_NUM_USAGES, 0),
+                call(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+            ]
+            for reset_call in reset_calls:
+                assert reset_call in mock_set_user_field.call_args_list, \
+                    f"Expected reset call {reset_call} not found"
+            
+            # Reset date should also be updated
+            assert any(call[0][0] == UJ_MITO_AI_LAST_RESET_DATE for call in mock_set_user_field.call_args_list), \
+                "Reset date should have been updated"
+        
+        # Check counter increment operations
+        if expected_ops["increment_chat"]:
+            expected_count = 1 if completion_count is None else completion_count + 1
+            chat_call = call(UJ_AI_MITO_API_NUM_USAGES, expected_count)
+            assert chat_call in mock_set_user_field.call_args_list, \
+                f"Expected chat completion increment call {chat_call} not found"
+        
+        if expected_ops["increment_autocomplete"]:
+            expected_count = 1 if autocomplete_count is None else autocomplete_count + 1
+            autocomplete_call = call(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, expected_count)
+            assert autocomplete_call in mock_set_user_field.call_args_list, \
+                f"Expected autocomplete increment call {autocomplete_call} not found"
+
+# Special edge cases for update_mito_server_quota
+def test_update_quota_exception_handling() -> None:
+    """Test that exceptions from set_user_field are properly propagated."""
+    
+    with (
+        patch("mito_ai.utils.server_limits.get_first_completion_date", return_value=CURRENT_DATE),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=CURRENT_DATE),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=10),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=20),
+        patch("mito_ai.utils.server_limits.set_user_field", side_effect=ValueError("Test error"))
+    ):
+        # Test with chat completion
+        with pytest.raises(ValueError) as exc_info:
+            update_mito_server_quota(MessageType.CHAT)
+        assert str(exc_info.value) == "Test error"
+        
+        # Test with inline completion
+        with pytest.raises(ValueError) as exc_info:
+            update_mito_server_quota(MessageType.INLINE_COMPLETION)
+        assert str(exc_info.value) == "Test error"
+
+def test_update_quota_future_reset_date() -> None:
+    """Test behavior when reset date is in the future."""
+    
+    mock_set_user_field = MagicMock()
+    
+    with (
+        patch("mito_ai.utils.server_limits.get_first_completion_date", return_value=CURRENT_DATE),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=FUTURE_DATE),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=10),
+        patch("mito_ai.utils.server_limits.set_user_field", mock_set_user_field)
+    ):
+        # Call the function
+        update_mito_server_quota(MessageType.CHAT)
+        
+        # Check that counters weren't reset
+        reset_calls = [
+            call(UJ_AI_MITO_API_NUM_USAGES, 0),
+            call(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+        ]
+        for reset_call in reset_calls:
+            assert reset_call not in mock_set_user_field.call_args_list, \
+                "Counters should not be reset when reset date is in the future"
+        
+        # Chat count should still be incremented
+        chat_call = call(UJ_AI_MITO_API_NUM_USAGES, 11)
+        assert chat_call in mock_set_user_field.call_args_list, \
+            "Chat completion count should still be incremented"
+
+def test_update_quota_chat_name_generation() -> None:
+    """
+    Test that when message_type is MessageType.CHAT_NAME_GENERATION, 
+    neither UJ_AI_MITO_API_NUM_USAGES nor UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES is incremented.
+    """
+    
+    # Mock set_user_field to track calls
+    mock_set_user_field = MagicMock()
+    
+    with (
+        patch("mito_ai.utils.server_limits.get_first_completion_date", return_value=CURRENT_DATE),
+        patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=CURRENT_DATE),
+        patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=10),
+        patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=20),
+        patch("mito_ai.utils.server_limits.set_user_field", mock_set_user_field)
+    ):
+        # Call the function with CHAT_NAME_GENERATION message type
+        update_mito_server_quota(MessageType.CHAT_NAME_GENERATION)
+        
+        # Verify that set_user_field was not called at all - no counters should be updated
+        mock_set_user_field.assert_not_called()
+        
+        # For comparison, call with a normal message type
+        update_mito_server_quota(MessageType.CHAT)
+        
+        # Now verify set_user_field was called to increment the chat counter
+        chat_call = call(UJ_AI_MITO_API_NUM_USAGES, 11)
+        assert chat_call in mock_set_user_field.call_args_list, \
+            "Chat completion count should be incremented for regular chat messages"

mito_ai/tests/settings/conftest.py

@@ -0,0 +1,26 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import pytest
+import shutil
+from mito_ai.tests.settings.test_settings_constants import (
+    SETTINGS_PATH,
+    BACKUP_SETTINGS_PATH,
+)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def backup_settings_json():
+    """Backup the settings.json file before tests and restore it after."""
+    if os.path.exists(SETTINGS_PATH):
+        # Create backup, so we can test with a clean settings.json file
+        shutil.move(SETTINGS_PATH, BACKUP_SETTINGS_PATH)
+
+    yield
+
+    # Cleanup after tests
+    if os.path.exists(SETTINGS_PATH):
+        os.remove(SETTINGS_PATH)
+    if os.path.exists(BACKUP_SETTINGS_PATH):
+        shutil.move(BACKUP_SETTINGS_PATH, SETTINGS_PATH)

mito_ai/tests/settings/settings_test.py

@@ -0,0 +1,70 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+
+SETTINGS_KEY = "super_secret_mode"
+SETTINGS_VALUE = "yes"
+
+from mito_ai.tests.conftest import TOKEN
+
+# --- PUT SETTINGS ---
+
+
+def test_put_settings_with_auth(jp_base_url: str) -> None:
+    response = requests.put(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+        headers={"Authorization": f"token {TOKEN}"},
+        json={"value": SETTINGS_VALUE},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["status"] == "updated"
+    assert response_json["key"] == SETTINGS_KEY
+    assert response_json["value"] == SETTINGS_VALUE
+
+
+def test_put_settings_with_no_auth(jp_base_url: str) -> None:
+    response = requests.put(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+        json={"value": SETTINGS_VALUE},
+    )
+
+    assert response.status_code == 403  # Forbidden
+
+
+def test_put_settings_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.put(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+        headers={"Authorization": f"token incorrect-token"},
+        json={"value": SETTINGS_VALUE},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# # --- GET SETTINGS ---
+
+
+def test_get_settings_with_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    assert response.json() == {"key": SETTINGS_KEY, "value": SETTINGS_VALUE}
+
+
+def test_get_settings_with_no_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_settings_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/settings/{SETTINGS_KEY}",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/settings/test_settings_constants.py

@@ -0,0 +1,9 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from typing import Final
+from mito_ai.utils.schema import MITO_FOLDER
+
+SETTINGS_PATH: Final[str] = os.path.join(MITO_FOLDER, "settings.json")
+BACKUP_SETTINGS_PATH: Final[str] = os.path.join(MITO_FOLDER, "settings_backup.json") 

mito_ai/tests/streamlit_conversion/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+