mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -0,0 +1,112 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+from unittest.mock import patch, MagicMock
+from mito_ai.streamlit_conversion.validate_streamlit_app import (
+    get_syntax_error,
+    get_runtime_errors,
+    validate_app
+)
+import pytest
+from mito_ai.path_utils import AbsoluteNotebookPath
+
+
+class TestGetSyntaxError:
+    """Test cases for get_syntax_error function"""
+
+    @pytest.mark.parametrize("code,expected_error,test_description", [
+        # Valid Python code should return no error
+        (
+                "import streamlit\nst.title('Hello World')",
+                None,
+                "valid Python code"
+        ),
+        # Invalid Python syntax should be caught
+        (
+                "import streamlit\nst.title('Hello World'",
+                "SyntaxError",
+                "invalid Python code"
+        ),
+        # Empty streamlit app is valid
+        (
+                "",
+                None,
+                "empty code"
+        ),
+    ])
+    def test_get_syntax_error(self, code, expected_error, test_description):
+        """Test syntax validation with various code inputs"""
+        error = get_syntax_error(code)
+
+        if expected_error is None:
+            assert error is None, f"Expected no error for {test_description}"
+        else:
+            assert error is not None, f"Expected error for {test_description}"
+            assert expected_error in error, f"Expected '{expected_error}' in error for {test_description}"
+
+class TestGetRuntimeErrors:
+    """Test cases for get_runtime_errors function"""
+
+    @pytest.mark.parametrize("app_code,expected_error", [
+        ("x = 5", None),
+        ("1/0", "division by zero"),
+        ("", None)
+    ])
+    def test_get_runtime_errors(self, app_code, expected_error):
+        """Test getting runtime errors"""
+        
+        absolute_path = AbsoluteNotebookPath('/notebook.ipynb')
+        errors = get_runtime_errors(app_code, absolute_path)
+        
+        if expected_error is None:
+            assert errors is None
+        else:
+            errors_str = str(errors)
+            assert expected_error in errors_str
+                
+    def test_get_runtime_errors_with_relative_path(self):
+        """Test getting runtime errors"""
+        
+        app_code ="""
+import streamlit as st
+import pandas as pd
+
+df=pd.read_csv('data.csv')
+"""
+        # Create a temporary csv file in the directory temp/data.csv
+        with tempfile.TemporaryDirectory() as temp_dir:
+            directory = 'app_directory'
+            csv_path = os.path.join(temp_dir, directory, "data.csv")
+            
+            os.makedirs(os.path.join(temp_dir, directory), exist_ok=True)
+            app_path = os.path.join(temp_dir, directory, "app.py")
+            
+            # Create the file if it doesn't exist
+            with open(csv_path, "w") as f:
+                f.write("name,age\nJohn,25\nJane,30")
+               
+            errors = get_runtime_errors(app_code, AbsoluteNotebookPath(app_path))
+            assert errors is None
+
+class TestValidateApp:
+    """Test cases for validate_app function"""
+
+    @pytest.mark.parametrize("app_code,expected_has_errors,expected_error_message", [
+        ("x=5", False, ""),
+        ("1/0", True, "division by zero"),
+        # Syntax errors are caught during runtime
+        ("", False, ""),
+    ])
+    def test_validate_app(self, app_code, expected_has_errors, expected_error_message):
+        """Test the validate_app function"""
+        errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
+        
+        has_errors = len(errors) > 0
+        assert has_errors == expected_has_errors
+        if expected_error_message:
+            errors_str = str(errors)
+            assert expected_error_message in errors_str
+        
+    

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -0,0 +1,118 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import os
+import tempfile
+from unittest.mock import patch, Mock, AsyncMock, MagicMock
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+from mito_ai.path_utils import AbsoluteNotebookPath
+
+
+class TestStreamlitPreviewHandler:
+    """Test cases for the StreamlitPreviewHandler."""
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "app_exists,force_recreate,streamlit_handler_called",
+        [
+            # Test case 1: App exists, not forcing recreate - should not call streamlit_handler
+            (True, False, False),
+            # Test case 2: App doesn't exist - should call streamlit_handler
+            (False, False, True),
+            # Test case 3: App exists but forcing recreate - should call streamlit_handler
+            (True, True, True),
+        ],
+        ids=[
+            "app_exists_no_force_recreate",
+            "app_does_not_exist_generates_new_one",
+            "app_exists_force_recreate",
+        ]
+    )
+    async def test_post_handler_app_generation(
+        self,
+        app_exists,
+        force_recreate,
+        streamlit_handler_called,
+    ):
+        """Test StreamlitPreviewHandler POST method with various scenarios."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            notebook_path = os.path.join(temp_dir, "test_notebook.ipynb")
+            notebook_id = "test_notebook_id"
+            # App file name is derived from notebook_id
+            app_file_name = f"{notebook_id}.py"
+            app_path = os.path.join(temp_dir, app_file_name)
+            
+            # Create notebook file
+            with open(notebook_path, "w") as f:
+                f.write('{"cells": []}')
+            
+            # Create app file if it should exist
+            if app_exists:
+                with open(app_path, "w") as f:
+                    f.write("import streamlit as st\nst.write('Hello World')")
+            
+            # Create a properly mocked Tornado application with required attributes
+            mock_application = MagicMock()
+            mock_application.ui_methods = {}
+            mock_application.ui_modules = {}
+            mock_application.settings = {}
+            
+            # Create a mock request with necessary tornado setup
+            mock_request = MagicMock()
+            mock_request.connection = MagicMock()
+            mock_request.connection.context = MagicMock()
+            
+            # Create handler instance  
+            handler = StreamlitPreviewHandler(
+                application=mock_application,
+                request=mock_request,
+            )
+            handler.initialize()
+            
+            # Mock authentication - set current_user to bypass @tornado.web.authenticated
+            handler.current_user = "test_user"  # type: ignore
+            
+            # Mock the finish method and other handler methods
+            finish_called = []
+            def mock_finish_func(response):
+                finish_called.append(response)
+            
+            # Mock streamlit_handler and preview manager
+            with patch.object(handler, 'get_json_body', return_value={
+                "notebook_path": notebook_path,
+                "notebook_id": notebook_id,
+                "force_recreate": force_recreate,
+                "edit_prompt": ""
+            }), \
+                 patch.object(handler, 'finish', side_effect=mock_finish_func), \
+                 patch.object(handler, 'set_status'), \
+                 patch('mito_ai.streamlit_preview.handlers.streamlit_handler', new_callable=AsyncMock) as mock_streamlit_handler, \
+                 patch.object(handler.preview_manager, 'start_streamlit_preview', return_value=8501) as mock_start_preview, \
+                 patch('mito_ai.streamlit_preview.handlers.log_streamlit_app_preview_success'):
+                
+                # Call the handler
+                await handler.post()  # type: ignore[misc]
+                
+                # Verify streamlit_handler was called or not called as expected
+                if streamlit_handler_called:
+                    assert mock_streamlit_handler.called
+                    # Verify it was called with the correct arguments
+                    call_args = mock_streamlit_handler.call_args
+                    assert call_args[0][0] == os.path.abspath(notebook_path)  # First argument should be the absolute notebook path
+                    assert call_args[0][1] == app_file_name  # Second argument should be the app file name
+                    assert call_args[0][2] == ""  # Third argument should be the edit_prompt
+                else:
+                    mock_streamlit_handler.assert_not_called()
+                
+                # Verify preview was started
+                mock_start_preview.assert_called_once()
+                
+                # Verify response was sent
+                assert len(finish_called) == 1
+                response = finish_called[0]
+                assert response["type"] == "success"
+                assert "port" in response
+                assert "id" in response
+
+    

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -0,0 +1,292 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import time
+import tempfile
+import os
+import shutil
+import subprocess
+import threading
+import requests
+import json
+from unittest.mock import Mock, patch, MagicMock
+from typing import Any
+
+from mito_ai.streamlit_preview.manager import (
+    StreamlitPreviewManager, 
+    PreviewProcess
+)
+
+
+class TestStreamlitPreviewManager:
+    """Test cases for StreamlitPreviewManager."""
+    
+    @pytest.fixture
+    def manager(self):
+        """Create a fresh manager instance for each test."""
+        return StreamlitPreviewManager()
+    
+    @pytest.fixture
+    def sample_app_code(self):
+        """Sample streamlit app code for testing."""
+        return """
+import streamlit as st
+
+st.title("Test App")
+st.write("Hello, World!")
+"""
+    
+    def test_init(self, manager):
+        """Test manager initialization."""
+        assert manager._previews == {}
+        assert isinstance(manager._lock, type(threading.Lock()))
+        assert manager.log is not None
+    
+    def test_get_free_port(self, manager):
+        """Test getting a free port."""
+        port = manager.get_free_port()
+        assert isinstance(port, int)
+        assert port > 0
+        assert port < 65536
+        
+        # Test that we get different ports
+        port2 = manager.get_free_port()
+        assert port != port2
+    
+    @pytest.mark.parametrize("app_code,preview_id,expected_success", [
+        ("import streamlit as st\nst.write('Hello')", "test_preview", True),
+        ("", "empty_preview", True),
+        ("import streamlit as st\n" * 1000 + "st.write('Large app')", "large_preview", True),
+    ])
+    def test_start_streamlit_preview_success_cases(self, manager, app_code, preview_id, expected_success):
+        """Test successful streamlit preview start with different app codes."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp:
+            
+            # Setup mocks
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Test
+            port = manager.start_streamlit_preview(app_directory, 'test-file-name.py',  preview_id)
+            
+            # Assertions
+            assert isinstance(port, int)
+            assert port > 0
+            
+            # Verify subprocess was called correctly
+            mock_popen.assert_called_once()
+            call_args = mock_popen.call_args
+            assert "streamlit" in call_args[0][0]
+            assert "run" in call_args[0][0]
+            assert "--server.headless" in call_args[0][0]
+            assert "--server.address" in call_args[0][0]
+            
+            # Cleanup
+            manager.stop_preview(preview_id)
+    
+    @pytest.mark.parametrize("exception_type,expected_message", [
+        (Exception("Temp dir creation failed"), "failed to start preview"),
+        (OSError("Permission denied"), "failed to start preview"),
+        (ValueError("Invalid argument"), "failed to start preview"),
+    ])
+    def test_start_streamlit_preview_exceptions(self, manager, sample_app_code, exception_type, expected_message):
+        """Test streamlit preview start with different exceptions."""
+        from mito_ai.utils.error_classes import StreamlitPreviewError
+        
+        with patch('subprocess.Popen', side_effect=exception_type):
+            app_directory = "/tmp/test_dir"
+            
+            with pytest.raises(StreamlitPreviewError) as exc_info:
+                manager.start_streamlit_preview(app_directory, 'test-file-name.py',  "test_preview")
+            
+            assert expected_message in str(exc_info.value).lower()
+    
+    @pytest.mark.parametrize("preview_id,expected_result", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_stop_preview_scenarios(self, manager, sample_app_code, preview_id, expected_result):
+        """Test stopping previews with different scenarios."""
+        if expected_result:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                app_directory = "/tmp/test_dir"
+                mock_mkdtemp.return_value = app_directory
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview(app_directory, 'test-file-name.py',  preview_id)
+    
+    @pytest.mark.parametrize("process_behavior,expected_kill_called", [
+        (subprocess.TimeoutExpired("cmd", 5), True),
+        (None, False),  # Normal termination
+    ])
+    def test_stop_preview_process_behaviors(self, manager, sample_app_code, process_behavior, expected_kill_called):
+        """Test stopping preview with different process behaviors."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks for start
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start a preview
+            manager.start_streamlit_preview(app_directory, 'test-file-name.py',  "test_preview")
+            
+            # Setup process behavior for stop
+            if process_behavior:
+                # Configure the mock to raise the exception when called with timeout
+                def wait_with_timeout(*args, **kwargs):
+                    if 'timeout' in kwargs:
+                        raise process_behavior
+                    return None
+                mock_proc.wait.side_effect = wait_with_timeout
+    
+    @pytest.mark.parametrize("preview_id,expected_found", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_get_preview_scenarios(self, manager, sample_app_code, preview_id, expected_found):
+        """Test getting previews with different scenarios."""
+        if expected_found:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                mock_mkdtemp.return_value = "/tmp/test_dir"
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview("/tmp/test_dir", 'test-file-name.py', preview_id)
+        
+        preview = manager.get_preview(preview_id)
+        
+        if expected_found:
+            assert preview is not None
+            assert isinstance(preview, PreviewProcess)
+            assert preview.port > 0
+            
+            # Cleanup
+            manager.stop_preview(preview_id)
+        else:
+            assert preview is None
+    
+    def test_preview_process_dataclass(self):
+        """Test PreviewProcess dataclass."""
+        proc = Mock()
+        port = 8080
+        
+        preview = PreviewProcess(
+            proc=proc,
+            port=port
+        )
+        
+        assert preview.proc == proc
+        assert preview.port == port
+
+    @pytest.mark.parametrize("num_previews", [1, 2, 3])
+    def test_concurrent_previews(self, manager, sample_app_code, num_previews):
+        """Test managing multiple concurrent previews."""
+        preview_ids = [f"preview_{i}" for i in range(num_previews)]
+        ports = []
+        
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks
+            mock_mkdtemp.return_value = "/tmp/test_dir"
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start multiple previews
+            for preview_id in preview_ids:
+                port = manager.start_streamlit_preview("/tmp/test_dir", 'test-file-name.py', preview_id)
+                ports.append(port)
+            
+            # Assertions
+            assert len(set(ports)) == num_previews  # All ports should be different
+            
+            # Check all previews exist
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is not None
+            
+            # Stop all previews
+            for preview_id in preview_ids:
+                assert manager.stop_preview(preview_id)
+            
+            # Verify they're gone
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is None
+

mito_ai/tests/test_constants.py

@@ -3,17 +3,45 @@
 
 from typing import Any
 import pytest
-from mito_ai.constants import ACTIVE_BASE_URL, MITO_PROD_BASE_URL, MITO_DEV_BASE_URL
+from mito_ai.constants import (
+    ACTIVE_BASE_URL, MITO_PROD_BASE_URL, MITO_DEV_BASE_URL,
+    MITO_STREAMLIT_DEV_BASE_URL, MITO_STREAMLIT_TEST_BASE_URL, ACTIVE_STREAMLIT_BASE_URL,
+    COGNITO_CONFIG_DEV, ACTIVE_COGNITO_CONFIG,
+)
 
 
 def test_prod_lambda_url() -> Any:
     """Make sure that the lambda urls are correct"""
-    assert MITO_PROD_BASE_URL == "https://yxwyadgaznhavqvgnbfuo2k6ca0jboku.lambda-url.us-east-1.on.aws"
+    assert MITO_PROD_BASE_URL.startswith("https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/")
     
 def test_dev_lambda_url() -> Any:
     """Make sure that the lambda urls are correct"""
-    assert MITO_DEV_BASE_URL == "https://x3rafympznv4abp7phos44gzgu0clbui.lambda-url.us-east-1.on.aws"
+    assert MITO_DEV_BASE_URL.startswith("https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/")
     
 def test_active_base_url() -> Any:
     """Make sure that the active base url is correct"""
     assert ACTIVE_BASE_URL == MITO_PROD_BASE_URL
+
+def test_devenv_streamlit_url() -> Any:
+    """Make sure that the streamlit urls are correct"""
+    assert MITO_STREAMLIT_DEV_BASE_URL == "https://fr12uvtfy5.execute-api.us-east-1.amazonaws.com"
+
+def test_testenv_streamlit_url() -> Any:
+    """Make sure that the streamlit urls are correct"""
+    assert MITO_STREAMLIT_TEST_BASE_URL == "https://iyual08t6d.execute-api.us-east-1.amazonaws.com"
+
+def test_streamlit_active_base_url() -> Any:
+    """Make sure that the active streamlit base url is correct"""
+    assert ACTIVE_STREAMLIT_BASE_URL == MITO_STREAMLIT_DEV_BASE_URL
+
+def test_cognito_config() -> Any:
+    """Make sure that the Cognito configuration is correct"""
+    expected_config = {
+        'TOKEN_ENDPOINT': 'https://mito-app-auth.auth.us-east-1.amazoncognito.com/oauth2/token',
+        'CLIENT_ID': '6ara3u3l8sss738hrhbq1qtiqf',
+        'CLIENT_SECRET': '',
+        'REDIRECT_URI': 'http://localhost:8888/lab'
+    }
+
+    assert COGNITO_CONFIG_DEV == expected_config
+    assert ACTIVE_COGNITO_CONFIG == COGNITO_CONFIG_DEV

mito_ai/tests/test_telemetry.py

@@ -0,0 +1,12 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from mito_ai.utils.telemetry_utils import PRINT_LOGS
+
+def test_print_logs_is_false():
+    """
+    Test to ensure that PRINT_LOGS is set to False.
+    """
+    assert not PRINT_LOGS, "PRINT_LOGS should be False by default."
+

mito_ai/tests/user/__init__.py

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