mito-ai 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -2,12 +2,12 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import pytest
-from unittest.mock import patch, AsyncMock
+from unittest.mock import patch, AsyncMock, MagicMock
 from mito_ai.streamlit_conversion.streamlit_agent_handler import (
     StreamlitCodeGeneration,
     streamlit_handler
 )
-from typing import cast
+from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
 
 # Add this line to enable async support
 pytest_plugins = ('pytest_asyncio',)
@@ -16,20 +16,24 @@ pytest_plugins = ('pytest_asyncio',)
 class TestStreamlitCodeGeneration:
     """Test cases for StreamlitCodeGeneration class"""
 
-    def test_init(self):
+    @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"""
+        # 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"]}]}
-        generator = StreamlitCodeGeneration(notebook_data)
-        
-        assert len(generator.messages) == 1
-        assert generator.messages[0]["role"] == "user"
-        # Access content properly as a list and cast to expected type
-        content_list = cast(list, generator.messages[0]["content"])
-        assert isinstance(content_list, list)
-        assert len(content_list) > 0
-        content_item = cast(dict, content_list[0])
-        assert content_item["type"] == "text"
-        assert "jupyter notebook content" in content_item["text"]
+        streamlit_code_handler = StreamlitCodeGeneration()
+        
+        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
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
@@ -48,9 +52,9 @@ class TestStreamlitCodeGeneration:
         mock_stream.return_value = mock_async_gen()
         
         notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
+        streamlit_code_handler = StreamlitCodeGeneration()
         
-        result = await generator.get_response_from_agent(generator.messages)
+        result = await streamlit_code_handler.generate_streamlit_code(notebook_data)
         
         assert result == expected_result
         mock_stream.assert_called_once()
@@ -63,28 +67,10 @@ class TestStreamlitCodeGeneration:
         mock_stream.side_effect = Exception("API Error")
         
         notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
+        streamlit_code_handler = StreamlitCodeGeneration()
         
         with pytest.raises(Exception, match="API Error"):
-            await generator.get_response_from_agent(generator.messages)
-
-    def test_add_agent_response_to_context(self):
-        """Test adding agent response to message history"""
-        notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
-        
-        initial_count = len(generator.messages)
-        generator.add_agent_response_to_context("Test response")
-        
-        assert len(generator.messages) == initial_count + 1
-        assert generator.messages[-1]["role"] == "assistant"
-        # Access content properly as a list and cast to expected type
-        content_list = cast(list, generator.messages[-1]["content"])
-        assert isinstance(content_list, list)
-        assert len(content_list) > 0
-        content_item = cast(dict, content_list[0])
-        assert content_item["type"] == "text"
-        assert content_item["text"] == "Test response"
+            await streamlit_code_handler.generate_streamlit_code(notebook_data)
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
@@ -99,22 +85,26 @@ class TestStreamlitCodeGeneration:
         mock_stream.return_value = mock_async_gen()
         
         notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
+        streamlit_code_handler = StreamlitCodeGeneration()
         
-        result = await generator.generate_streamlit_code()
+        result = await streamlit_code_handler.generate_streamlit_code(notebook_data)
         
         expected_code = "import streamlit\nst.title('Hello')\n"
         assert result == expected_code
-        
-        # Check that response was added to context
-        assert len(generator.messages) == 2
-        assert generator.messages[-1]["role"] == "assistant"
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
     async def test_correct_error_in_generation_success(self, mock_stream):
         """Test successful error correction"""
-        mock_response = "Here's the corrected code:\n```python\nimport streamlit\nst.title('Fixed')\n```"
+        mock_response = """```unified_diff
+--- a/app.py
++++ b/app.py
+@@ -1,1 +1,1 @@
+-import streamlit
+-st.title('Test')
++import streamlit
++st.title('Fixed')
+```"""
         async def mock_async_gen():
             for item in [mock_response]:
                 yield item
@@ -122,29 +112,23 @@ class TestStreamlitCodeGeneration:
         mock_stream.return_value = mock_async_gen()
 
         notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
+        streamlit_code_handler = StreamlitCodeGeneration()
 
-        result = await generator.correct_error_in_generation("ImportError: No module named 'pandas'")
+        result = await streamlit_code_handler.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
 
-        # Check that error message and response were added to context
-        assert len(generator.messages) == 3
-        assert generator.messages[-2]["role"] == "user"
-        assert generator.messages[-1]["role"] == "assistant"
-
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.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")
         
-        notebook_data: dict = {"cells": []}
-        generator = StreamlitCodeGeneration(notebook_data)
+        streamlit_code_handler = StreamlitCodeGeneration()
         
         with pytest.raises(Exception, match="API Error"):
-            await generator.correct_error_in_generation("Some error")
+            await streamlit_code_handler.correct_error_in_generation("Some error", "import streamlit\nst.title('Test')")
 
 
 class TestStreamlitHandler:
@@ -153,9 +137,10 @@ 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.streamlit_code_validator')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    async def test_streamlit_handler_success(self, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+    @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):
         """Test successful streamlit handler execution"""
         # Mock notebook parsing
         mock_notebook_data: dict = {"cells": [{"cell_type": "code", "source": ["import pandas"]}]}
@@ -170,43 +155,47 @@ class TestStreamlitHandler:
         mock_validator.return_value = (False, "")
         
         # Mock file creation
-        mock_create_file.return_value = (True, "File created successfully")
+        mock_create_file.return_value = (True, "/path/to/app.py", "File created successfully")
+        
+        # Mock clean directory check (no-op)
+        mock_clean_directory.return_value = None
         
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        result = await streamlit_handler("/path/to/notebook.ipynb")
         
         assert result[0] is True
-        assert "File created successfully" in result[1]
+        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_notebook_data)
+        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')")
-        mock_create_file.assert_called_once_with("/path/to/app", "import streamlit\nst.title('Test')")
+        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')")
 
     @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.streamlit_code_validator')
+    @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):
         """Test streamlit handler when max retries are exceeded"""
         # Mock notebook parsing
         mock_notebook_data: 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 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")
         
-        # Mock validation (always errors)
-        mock_validator.return_value = (True, "Persistent error")
-        
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
-        
+        # Verify the result indicates failure
         assert result[0] is False
-        assert "Error generating streamlit code by agent" in result[1]
+        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
@@ -214,9 +203,10 @@ 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.streamlit_code_validator')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    async def test_streamlit_handler_file_creation_failure(self, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+    @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):
         """Test streamlit handler when file creation fails"""
         # Mock notebook parsing
         mock_notebook_data: dict = {"cells": []}
@@ -231,26 +221,34 @@ class TestStreamlitHandler:
         mock_validator.return_value = (False, "")
         
         # Mock file creation failure
-        mock_create_file.return_value = (False, "Permission denied")
+        mock_create_file.return_value = (False, None, "Permission denied")
+        
+        # Mock clean directory check (no-op)
+        mock_clean_directory.return_value = None
         
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        result = await streamlit_handler("/path/to/notebook.ipynb")
         
         assert result[0] is False
-        assert "Permission denied" in result[1]
+        assert "Permission denied" in result[2]
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    async def test_streamlit_handler_parse_notebook_exception(self, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
+    async def test_streamlit_handler_parse_notebook_exception(self, mock_clean_directory, mock_parse):
         """Test streamlit handler when notebook parsing fails"""
+        # Mock clean directory check (no-op)
+        mock_clean_directory.return_value = None
+        
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+            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')
-    async def test_streamlit_handler_generation_exception(self, mock_generator_class, mock_parse):
+    @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):
         """Test streamlit handler when code generation fails"""
         # Mock notebook parsing
         mock_notebook_data: dict = {"cells": []}
@@ -261,5 +259,95 @@ class TestStreamlitHandler:
         mock_generator.generate_streamlit_code.side_effect = Exception("Generation failed")
         mock_generator_class.return_value = mock_generator
         
+        # 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", "/path/to/app")
+            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()
+
+
+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

@@ -13,7 +13,6 @@ from mito_ai.streamlit_conversion.streamlit_utils import (
 )
 from typing import Dict, Any
 
-
 class TestExtractCodeBlocks:
     """Test cases for extract_code_blocks function"""
 
@@ -52,7 +51,7 @@ class TestCreateAppFile:
         file_path = str(tmp_path)
         code = "import streamlit\nst.title('Test')"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is True
         assert "Successfully created" in message
@@ -70,7 +69,7 @@ class TestCreateAppFile:
         file_path = "/nonexistent/path/that/should/fail"
         code = "import streamlit"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is False
         assert "Error creating file" in message
@@ -81,7 +80,7 @@ class TestCreateAppFile:
         file_path = "/tmp/test"
         code = "import streamlit"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is False
         assert "Unexpected error" in message
@@ -91,7 +90,7 @@ class TestCreateAppFile:
         file_path = str(tmp_path)
         code = ""
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is True
         assert "Successfully created" in message

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -0,0 +1,119 @@
+# 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 (
+    StreamlitValidator,
+    validate_app
+)
+import pytest
+
+
+class TestStreamlitValidator:
+    """Test cases for StreamlitValidator class"""
+
+    @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_validate_syntax(self, code, expected_error, test_description):
+        """Test syntax validation with various code inputs"""
+        validator = StreamlitValidator()
+        
+        error = validator.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}"
+        
+    @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"""
+        validator = StreamlitValidator()
+        
+        errors = validator.get_runtime_errors(app_code, '/app.py')
+        
+        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")
+               
+            validator = StreamlitValidator() 
+            errors = validator.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()
+
+
+    @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
+        ("x=5", False, ""),
+        ("1/0", True, "division by zero"),
+        ("print('Hello World'", True, "SyntaxError"),
+        ("", False, ""),
+    ])
+    def test_streamlit_code_validator(self, app_code, expected_has_validation_error, expected_error_message):
+
+        has_validation_error, errors = validate_app(app_code, '/app.py')
+        
+        assert has_validation_error == expected_has_validation_error
+        assert expected_error_message in str(errors)
+        
+