mito-ai 0.1.44__py3-none-any.whl → 0.1.46__py3-none-any.whl

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -1,10 +1,15 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
+from typing import List
+from anthropic.types import MessageParam
 import pytest
+import os
 from unittest.mock import patch, AsyncMock, MagicMock
 from mito_ai.streamlit_conversion.streamlit_agent_handler import (
-    StreamlitCodeGeneration,
+    get_response_from_agent,
+    generate_new_streamlit_code,
+    correct_error_in_generation,
     streamlit_handler
 )
 from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
@@ -13,36 +18,34 @@ from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
 pytest_plugins = ('pytest_asyncio',)
 
 
-class TestStreamlitCodeGeneration:
-    """Test cases for StreamlitCodeGeneration class"""
+class TestGetResponseFromAgent:
+    """Test cases for get_response_from_agent function"""
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
-    async def test_init(self, mock_stream):
-        """Test StreamlitCodeGeneration initialization"""
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
+    async def test_get_response_from_agent_success(self, mock_stream):
+        """Test get_response_from_agent with successful response"""
         # Mock the async generator
         async def mock_async_gen():
             yield "Here's your code:\n```python\nimport streamlit\nst.title('Test')\n```"
 
         mock_stream.return_value = mock_async_gen()
         
-        notebook_data: dict = {"cells": [{"cell_type": "code", "source": ["import pandas"]}]}
-        streamlit_code_handler = StreamlitCodeGeneration()
+        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
+        response = await get_response_from_agent(messages)
         
-        streamlit_code = await streamlit_code_handler.generate_streamlit_code(notebook_data)
-        
-        assert streamlit_code is not None
-        assert len(streamlit_code) > 0
-        assert "import streamlit" in streamlit_code
+        assert response is not None
+        assert len(response) > 0
+        assert "import streamlit" in response
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
     @pytest.mark.parametrize("mock_items,expected_result", [
         (["Hello", " World", "!"], "Hello World!"),
         ([], ""),
         (["Here's your code: import streamlit"], "Here's your code: import streamlit")
     ])
-    async def test_get_response_from_agent(self, mock_stream, mock_items, expected_result):
+    async def test_get_response_from_agent_parametrized(self, mock_stream, mock_items, expected_result):
         """Test response from agent with different scenarios"""
         # Mock the async generator
         async def mock_async_gen():
@@ -51,30 +54,31 @@ class TestStreamlitCodeGeneration:
 
         mock_stream.return_value = mock_async_gen()
         
-        notebook_data: dict = {"cells": []}
-        streamlit_code_handler = StreamlitCodeGeneration()
-        
-        result = await streamlit_code_handler.generate_streamlit_code(notebook_data)
+        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
+        result = await get_response_from_agent(messages)
         
         assert result == expected_result
         mock_stream.assert_called_once()
         
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
     async def test_get_response_from_agent_exception(self, mock_stream):
         """Test exception handling in get_response_from_agent"""
         mock_stream.side_effect = Exception("API Error")
         
-        notebook_data: dict = {"cells": []}
-        streamlit_code_handler = StreamlitCodeGeneration()
+        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
         
         with pytest.raises(Exception, match="API Error"):
-            await streamlit_code_handler.generate_streamlit_code(notebook_data)
+            await get_response_from_agent(messages)
+
+
+class TestGenerateStreamlitCode:
+    """Test cases for generate_new_streamlit_code function"""
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
-    async def test_generate_streamlit_code_success(self, mock_stream):
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
+    async def test_generate_new_streamlit_code_success(self, mock_stream):
         """Test successful streamlit code generation"""
         mock_response = "Here's your code:\n```python\nimport streamlit\nst.title('Hello')\n```"
 
@@ -84,16 +88,18 @@ class TestStreamlitCodeGeneration:
 
         mock_stream.return_value = mock_async_gen()
         
-        notebook_data: dict = {"cells": []}
-        streamlit_code_handler = StreamlitCodeGeneration()
-        
-        result = await streamlit_code_handler.generate_streamlit_code(notebook_data)
+        notebook_data: List[dict] = [{"cells": []}]
+        result = await generate_new_streamlit_code(notebook_data)
         
         expected_code = "import streamlit\nst.title('Hello')\n"
         assert result == expected_code
 
+
+class TestCorrectErrorInGeneration:
+    """Test cases for correct_error_in_generation function"""
+
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
     async def test_correct_error_in_generation_success(self, mock_stream):
         """Test successful error correction"""
         mock_response = """```unified_diff
@@ -111,24 +117,19 @@ class TestStreamlitCodeGeneration:
 
         mock_stream.return_value = mock_async_gen()
 
-        notebook_data: dict = {"cells": []}
-        streamlit_code_handler = StreamlitCodeGeneration()
-
-        result = await streamlit_code_handler.correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')")
+        result = await correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')")
 
         expected_code = "import streamlit\nst.title('Fixed')\n"
         assert result == expected_code
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
     async def test_correct_error_in_generation_exception(self, mock_stream):
         """Test exception handling in error correction"""
         mock_stream.side_effect = Exception("API Error")
         
-        streamlit_code_handler = StreamlitCodeGeneration()
-        
         with pytest.raises(Exception, match="API Error"):
-            await streamlit_code_handler.correct_error_in_generation("Some error", "import streamlit\nst.title('Test')")
+            await correct_error_in_generation("Some error", "import streamlit\nst.title('Test')")
 
 
 class TestStreamlitHandler:
@@ -136,20 +137,18 @@ class TestStreamlitHandler:
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_success(self, mock_clean_directory, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+    async def test_streamlit_handler_success(self, mock_clean_directory, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test successful streamlit handler execution"""
         # Mock notebook parsing
-        mock_notebook_data: dict = {"cells": [{"cell_type": "code", "source": ["import pandas"]}]}
+        mock_notebook_data: List[dict] = [{"cells": [{"cell_type": "code", "source": ["import pandas"]}]}]
         mock_parse.return_value = mock_notebook_data
         
         # Mock code generation
-        mock_generator = AsyncMock()
-        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
-        mock_generator_class.return_value = mock_generator
+        mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
         # Mock validation (no errors)
         mock_validator.return_value = (False, "")
@@ -160,62 +159,62 @@ class TestStreamlitHandler:
         # Mock clean directory check (no-op)
         mock_clean_directory.return_value = None
         
-        result = await streamlit_handler("/path/to/notebook.ipynb")
+        # Use a relative path that will work cross-platform
+        notebook_path = "notebook.ipynb"
+        result = await streamlit_handler(notebook_path)
         
         assert result[0] is True
         assert "File created successfully" in result[2]
         
         # Verify calls
-        mock_parse.assert_called_once_with("/path/to/notebook.ipynb")
-        mock_generator_class.assert_called_once_with()
-        mock_generator.generate_streamlit_code.assert_called_once()
-        mock_validator.assert_called_once_with("import streamlit\nst.title('Test')", "/path/to/notebook.ipynb")
-        mock_create_file.assert_called_once_with("/path/to", "import streamlit\nst.title('Test')")
+        mock_parse.assert_called_once_with(notebook_path)
+        mock_generate_code.assert_called_once_with(mock_notebook_data)
+        mock_validator.assert_called_once_with("import streamlit\nst.title('Test')", notebook_path)
+        # get_app_directory converts relative paths to absolute, so expect the absolute path directory
+        expected_app_dir = os.path.dirname(os.path.abspath(notebook_path))
+        mock_create_file.assert_called_once_with(expected_app_dir, "import streamlit\nst.title('Test')")
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.correct_error_in_generation')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
-    async def test_streamlit_handler_max_retries_exceeded(self, mock_validator, mock_generator_class, mock_parse):
+    async def test_streamlit_handler_max_retries_exceeded(self, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
         """Test streamlit handler when max retries are exceeded"""
         # Mock notebook parsing
-        mock_notebook_data: dict = {"cells": []}
+        mock_notebook_data: List[dict] = [{"cells": []}]
         mock_parse.return_value = mock_notebook_data
     
         # Mock code generation
-        mock_generator = AsyncMock()
-        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
-        mock_generator.correct_error_in_generation.return_value = "import streamlit\nst.title('Fixed')"
-        mock_generator_class.return_value = mock_generator
+        mock_generate_code.return_value = "import streamlit\nst.title('Test')"
+        mock_correct_error.return_value = "import streamlit\nst.title('Fixed')"
     
         # Mock validation (always errors) - Return list of errors as expected by validate_app
         mock_validator.return_value = (True, ["Persistent error"])
     
-        result = await streamlit_handler("/notebook.ipynb")
+        result = await streamlit_handler("notebook.ipynb")
         
         # Verify the result indicates failure
         assert result[0] is False
         assert "Error generating streamlit code by agent" in result[2]
         
         # Verify that error correction was called 5 times (max retries)
-        assert mock_generator.correct_error_in_generation.call_count == 5
+        assert mock_correct_error.call_count == 5
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_file_creation_failure(self, mock_clean_directory, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+    async def test_streamlit_handler_file_creation_failure(self, mock_clean_directory, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test streamlit handler when file creation fails"""
         # Mock notebook parsing
-        mock_notebook_data: dict = {"cells": []}
+        mock_notebook_data: List[dict] = [{"cells": []}]
         mock_parse.return_value = mock_notebook_data
         
         # Mock code generation
-        mock_generator = AsyncMock()
-        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
-        mock_generator_class.return_value = mock_generator
+        mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
         # Mock validation (no errors)
         mock_validator.return_value = (False, "")
@@ -226,7 +225,7 @@ class TestStreamlitHandler:
         # Mock clean directory check (no-op)
         mock_clean_directory.return_value = None
         
-        result = await streamlit_handler("/path/to/notebook.ipynb")
+        result = await streamlit_handler("notebook.ipynb")
         
         assert result[0] is False
         assert "Permission denied" in result[2]
@@ -242,112 +241,26 @@ class TestStreamlitHandler:
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler("/path/to/notebook.ipynb")
+            await streamlit_handler("notebook.ipynb")
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_generation_exception(self, mock_clean_directory, mock_generator_class, mock_parse):
+    async def test_streamlit_handler_generation_exception(self, mock_clean_directory, mock_generate_code, mock_parse):
         """Test streamlit handler when code generation fails"""
         # Mock notebook parsing
-        mock_notebook_data: dict = {"cells": []}
+        mock_notebook_data: List[dict] = [{"cells": []}]
         mock_parse.return_value = mock_notebook_data
         
         # Mock code generation failure
-        mock_generator = AsyncMock()
-        mock_generator.generate_streamlit_code.side_effect = Exception("Generation failed")
-        mock_generator_class.return_value = mock_generator
+        mock_generate_code.side_effect = Exception("Generation failed")
         
         # Mock clean directory check (no-op)
         mock_clean_directory.return_value = None
         
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler("/path/to/notebook.ipynb")
-
-    @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_too_many_files_in_directory(self, mock_clean_directory, mock_create_file, mock_validator, mock_generator_class, mock_parse):
-        """Test streamlit handler when there are too many files in the directory"""
-        # Mock clean directory check to raise ValueError (simulating >10 files)
-        mock_clean_directory.side_effect = ValueError("Too many files in directory: 10 allowed but 15 present. Create a new directory and retry")
-        
-        # The function should raise the ValueError before any other processing
-        with pytest.raises(ValueError, match="Too many files in directory: 10 allowed but 15 present. Create a new directory and retry"):
-            await streamlit_handler("/path/to/notebook.ipynb")
-        
-        # Verify that clean_directory_check was called
-        mock_clean_directory.assert_called_once_with("/path/to/notebook.ipynb")
-        
-        # Verify that no other functions were called since the error occurred early
-        mock_parse.assert_not_called()
-        mock_generator_class.assert_not_called()
-        mock_validator.assert_not_called()
-        mock_create_file.assert_not_called()
+            await streamlit_handler("notebook.ipynb")
 
 
-class TestCleanDirectoryCheck:
-    """Test cases for clean_directory_check function"""
 
-    @patch('mito_ai.streamlit_conversion.streamlit_utils.Path')
-    def test_clean_directory_check_under_limit(self, mock_path):
-        """Test clean_directory_check when directory has 10 or fewer files"""
-        # Mock the Path class and its methods
-        mock_path_instance = mock_path.return_value
-        mock_path_instance.resolve.return_value = mock_path_instance
-        mock_path_instance.parent = mock_path_instance
-        
-        # Mock directory existence check
-        mock_path_instance.exists.return_value = True
-        
-        # Mock directory contents with 8 files
-        mock_files = []
-        for i in range(8):
-            mock_file = MagicMock()
-            mock_file.is_file.return_value = True
-            mock_files.append(mock_file)
-        
-        mock_path_instance.iterdir.return_value = mock_files
-        
-        # Should not raise any exception
-        clean_directory_check('/path/to/notebook.ipynb')
-        
-        # Verify calls
-        mock_path.assert_called_once_with('/path/to/notebook.ipynb')
-        mock_path_instance.resolve.assert_called_once()
-        mock_path_instance.exists.assert_called_once()
-        mock_path_instance.iterdir.assert_called_once()
-
-    @patch('mito_ai.streamlit_conversion.streamlit_utils.Path')
-    def test_clean_directory_check_over_limit(self, mock_path):
-        """Test clean_directory_check when directory has more than 10 files"""
-        # Mock the Path class and its methods
-        mock_path_instance = mock_path.return_value
-        mock_path_instance.resolve.return_value = mock_path_instance
-        mock_path_instance.parent = mock_path_instance
-        
-        # Mock directory existence check
-        mock_path_instance.exists.return_value = True
-        
-        # Mock directory contents with 15 files
-        mock_files = []
-        for i in range(15):
-            mock_file = MagicMock()
-            mock_file.is_file.return_value = True
-            mock_files.append(mock_file)
-        
-        mock_path_instance.iterdir.return_value = mock_files
-        
-        # Should raise ValueError
-        with pytest.raises(ValueError, match="Too many files in directory: 10 allowed but 15 present. Create a new directory and retry"):
-            clean_directory_check('/path/to/notebook.ipynb')
-        
-        # Verify calls
-        mock_path.assert_called_once_with('/path/to/notebook.ipynb')
-        mock_path_instance.resolve.assert_called_once()
-        mock_path_instance.exists.assert_called_once()
-        mock_path_instance.iterdir.assert_called_once()

mito_ai/tests/streamlit_conversion/test_streamlit_utils.py

@@ -131,15 +131,15 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
         
         # Check that only cell_type and source are preserved
-        assert len(result['cells']) == 2
-        assert result['cells'][0]['cell_type'] == 'code'
-        assert result['cells'][0]['source'] == ["import pandas as pd\n", "df = pd.DataFrame()\n"]
-        assert 'metadata' not in result['cells'][0]
-        assert 'execution_count' not in result['cells'][0]
+        assert len(result) == 2
+        assert result[0]['cell_type'] == 'code'
+        assert result[0]['source'] == ["import pandas as pd\n", "df = pd.DataFrame()\n"]
+        assert 'metadata' not in result[0]
+        assert 'execution_count' not in result[0]
         
-        assert result['cells'][1]['cell_type'] == 'markdown'
-        assert result['cells'][1]['source'] == ["# Title\n", "Some text\n"]
-        assert 'metadata' not in result['cells'][1]
+        assert result[1]['cell_type'] == 'markdown'
+        assert result[1]['source'] == ["# Title\n", "Some text\n"]
+        assert 'metadata' not in result[1]
 
     def test_parse_notebook_file_not_found(self):
         """Test parsing non-existent notebook file"""
@@ -171,15 +171,15 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         
         result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
         
-        assert len(result['cells']) == 3
-        assert result['cells'][0]['cell_type'] == 'code'
-        assert result['cells'][0]['source'] == []  # Default empty list
+        assert len(result) == 3
+        assert result[0]['cell_type'] == 'code'
+        assert result[0]['source'] == []  # Default empty list
         
-        assert result['cells'][1]['cell_type'] == ''  # Default empty string
-        assert result['cells'][1]['source'] == ["some text"]
+        assert result[1]['cell_type'] == ''  # Default empty string
+        assert result[1]['source'] == ["some text"]
         
-        assert result['cells'][2]['cell_type'] == 'markdown'
-        assert result['cells'][2]['source'] == ["# Title"]
+        assert result[2]['cell_type'] == 'markdown'
+        assert result[2]['source'] == ["# Title"]
 
     def test_parse_empty_notebook(self, tmp_path):
         """Test parsing notebook with empty cells list"""
@@ -193,4 +193,4 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         
         result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
         
-        assert result['cells'] == []
+        assert result == []

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -5,14 +5,16 @@ import os
 import tempfile
 from unittest.mock import patch, MagicMock
 from mito_ai.streamlit_conversion.validate_streamlit_app import (
-    StreamlitValidator,
+    get_syntax_error,
+    get_runtime_errors,
+    check_for_errors,
     validate_app
 )
 import pytest
 
 
-class TestStreamlitValidator:
-    """Test cases for StreamlitValidator class"""
+class TestGetSyntaxError:
+    """Test cases for get_syntax_error function"""
 
     @pytest.mark.parametrize("code,expected_error,test_description", [
         # Valid Python code should return no error
@@ -34,11 +36,9 @@ class TestStreamlitValidator:
             "empty code"
         ),
     ])
-    def test_validate_syntax(self, code, expected_error, test_description):
+    def test_get_syntax_error(self, code, expected_error, test_description):
         """Test syntax validation with various code inputs"""
-        validator = StreamlitValidator()
-        
-        error = validator.get_syntax_error(code)
+        error = get_syntax_error(code)
         
         if expected_error is None:
             assert error is None, f"Expected no error for {test_description}"
@@ -46,6 +46,9 @@ class TestStreamlitValidator:
             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"),
@@ -53,9 +56,7 @@ class TestStreamlitValidator:
     ])
     def test_get_runtime_errors(self, app_code, expected_error):
         """Test getting runtime errors"""
-        validator = StreamlitValidator()
-        
-        errors = validator.get_runtime_errors(app_code, '/app.py')
+        errors = get_runtime_errors(app_code, '/app.py')
         
         if expected_error is None:
             assert errors is None
@@ -84,24 +85,11 @@ df=pd.read_csv('data.csv')
             with open(csv_path, "w") as f:
                 f.write("name,age\nJohn,25\nJane,30")
                
-            validator = StreamlitValidator() 
-            errors = validator.get_runtime_errors(app_code, app_path)
+            errors = get_runtime_errors(app_code, app_path)
             assert errors is None
-            
-            
-    @patch('subprocess.Popen')
-    def test_cleanup_with_process(self, mock_popen):
-        """Test cleanup with running process"""
-        validator = StreamlitValidator()
-        validator.temp_dir = "/tmp/test_dir"
-        
-        # Mock directory exists
-        with patch('os.path.exists', return_value=True):
-            with patch('shutil.rmtree') as mock_rmtree:
-                validator.cleanup()
-                
-                mock_rmtree.assert_called_once()
 
+class TestValidateApp:
+    """Test cases for validate_app function"""
 
     @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
         ("x=5", False, ""),
@@ -109,8 +97,8 @@ df=pd.read_csv('data.csv')
         ("print('Hello World'", True, "SyntaxError"),
         ("", False, ""),
     ])
-    def test_streamlit_code_validator(self, app_code, expected_has_validation_error, expected_error_message):
-
+    def test_validate_app(self, app_code, expected_has_validation_error, expected_error_message):
+        """Test the validate_app function"""
         has_validation_error, errors = validate_app(app_code, '/app.py')
         
         assert has_validation_error == expected_has_validation_error

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -70,7 +70,7 @@ class TestEnsureAppExists:
                     if streamlit_handler_return is not None:
                         mock_streamlit_handler.return_value = streamlit_handler_return
                     
-                    success, error_msg = await ensure_app_exists(notebook_path)
+                    success, error_msg = await ensure_app_exists(notebook_path, False, "")
                     
                     # Assertions
                     assert success == expected_success
@@ -81,7 +81,7 @@ class TestEnsureAppExists:
                     
                     # Verify streamlit_handler was called or not called as expected
                     if streamlit_handler_called:
-                        mock_streamlit_handler.assert_called_once_with(notebook_path)
+                        mock_streamlit_handler.assert_called_once_with(notebook_path, "")
                     else:
                         mock_streamlit_handler.assert_not_called()
 

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.