mito-ai 0.1.45__py3-none-any.whl → 0.1.47__py3-none-any.whl

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -1,48 +1,51 @@
 # 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
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
 
 # Add this line to enable async support
 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,26 +88,26 @@ 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
---- a/app.py
-+++ b/app.py
-@@ -1,1 +1,1 @@
--import streamlit
--st.title('Test')
-+import streamlit
-+st.title('Fixed')
+        mock_response = """```search_replace
+>>>>>>> SEARCH
+st.title('Test')
+=======
+st.title('Fixed')
+<<<<<<< REPLACE
 ```"""
         async def mock_async_gen():
             for item in [mock_response]:
@@ -111,24 +115,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')\n")
+        
         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,134 +135,108 @@ 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_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, "")
         
-        # Mock file creation
-        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")
+        # Use a relative path that will work cross-platform
+        notebook_path = AbsoluteNotebookPath("absolute/path/to/notebook.ipynb")
         
-        assert result[0] is True
-        assert "File created successfully" in result[2]
+        # Construct the expected app path using the same method as the production code
+        app_directory = get_absolute_notebook_dir_path(notebook_path)
+        expected_app_path = get_absolute_app_path(app_directory)
+        await streamlit_handler(notebook_path)
         
         # 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)
+        mock_create_file.assert_called_once_with(expected_app_path, "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")
-        
-        # Verify the result indicates failure
-        assert result[0] is False
-        assert "Error generating streamlit code by agent" in result[2]
+        # Now it should raise an exception instead of returning a tuple
+        with pytest.raises(Exception):
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
         
         # 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_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, "")
         
-        # Mock file creation failure
-        mock_create_file.return_value = (False, None, "Permission denied")
-        
-        # Mock clean directory check (no-op)
-        mock_clean_directory.return_value = None
+        # Mock file creation failure - now it should raise an exception
+        mock_create_file.side_effect = Exception("Permission denied")
         
-        result = await streamlit_handler("/path/to/notebook.ipynb")
-        
-        assert result[0] is False
-        assert "Permission denied" in result[2]
+        # Now it should raise an exception instead of returning a tuple
+        with pytest.raises(Exception, match="Permission denied"):
+            await streamlit_handler(AbsoluteNotebookPath("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.clean_directory_check')
-    async def test_streamlit_handler_parse_notebook_exception(self, mock_clean_directory, mock_parse):
+    async def test_streamlit_handler_parse_notebook_exception(self, 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")
+            await streamlit_handler(AbsoluteNotebookPath("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.clean_directory_check')
-    async def test_streamlit_handler_generation_exception(self, mock_clean_directory, mock_generator_class, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
+    async def test_streamlit_handler_generation_exception(self, 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 clean directory check (no-op)
-        mock_clean_directory.return_value = None
+        mock_generate_code.side_effect = Exception("Generation failed")
         
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler("/path/to/notebook.ipynb")
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
 
 
 

mito_ai/tests/streamlit_conversion/test_streamlit_utils.py

@@ -11,6 +11,7 @@ from mito_ai.streamlit_conversion.streamlit_utils import (
     create_app_file,
     parse_jupyter_notebook_to_extract_required_content
 )
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookDirPath, AbsoluteNotebookPath, get_absolute_notebook_path
 from typing import Dict, Any
 
 class TestExtractCodeBlocks:
@@ -48,55 +49,47 @@ class TestCreateAppFile:
 
     def test_create_app_file_success(self, tmp_path):
         """Test successful file creation"""
-        file_path = str(tmp_path)
+        app_path = os.path.join(str(tmp_path), "app.py")
         code = "import streamlit\nst.title('Test')"
         
-        success, app_path, message = create_app_file(file_path, code)
+        create_app_file(AbsoluteAppPath(app_path), code)
         
-        assert success is True
-        assert "Successfully created" in message
+        assert app_path is not None
+        assert os.path.exists(app_path)
         
         # Verify file was created with correct content
-        app_file_path = os.path.join(file_path, "app.py")
-        assert os.path.exists(app_file_path)
-        
-        with open(app_file_path, 'r') as f:
+        with open(app_path, 'r') as f:
             content = f.read()
         assert content == code
 
     def test_create_app_file_io_error(self):
         """Test file creation with IO error"""
-        file_path = "/nonexistent/path/that/should/fail"
+        file_path = AbsoluteAppPath("/nonexistent/path/that/should/fail")
         code = "import streamlit"
         
-        success, app_path, message = create_app_file(file_path, code)
-        
-        assert success is False
-        assert "Error creating file" in message
+        with pytest.raises(Exception):
+            create_app_file(file_path, code)
 
     @patch('builtins.open', side_effect=Exception("Unexpected error"))
     def test_create_app_file_unexpected_error(self, mock_open):
         """Test file creation with unexpected error"""
-        file_path = "/tmp/test"
+        app_path = AbsoluteAppPath("/tmp/test")
         code = "import streamlit"
         
-        success, app_path, message = create_app_file(file_path, code)
-        
-        assert success is False
-        assert "Unexpected error" in message
+        with pytest.raises(Exception, match="Unexpected error"):
+            create_app_file(app_path, code)
 
     def test_create_app_file_empty_code(self, tmp_path):
         """Test creating file with empty code"""
-        file_path = str(tmp_path)
+        app_path = AbsoluteAppPath(os.path.join(str(tmp_path), "app.py"))
         code = ""
         
-        success, app_path, message = create_app_file(file_path, code)
+        create_app_file(app_path, code)
         
-        assert success is True
-        assert "Successfully created" in message
+        assert app_path is not None
+        assert os.path.exists(app_path)
         
-        app_file_path = os.path.join(file_path, "app.py")
-        with open(app_file_path, 'r') as f:
+        with open(app_path, 'r') as f:
             content = f.read()
         assert content == ""
 
@@ -128,23 +121,25 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_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"""
-        with pytest.raises(FileNotFoundError, match="Notebook file not found"):
-            parse_jupyter_notebook_to_extract_required_content("/nonexistent/path/notebook.ipynb")
+        from mito_ai.utils.error_classes import StreamlitConversionError
+        with pytest.raises(StreamlitConversionError, match="Notebook file not found"):
+            parse_jupyter_notebook_to_extract_required_content(AbsoluteNotebookPath("/nonexistent/path/notebook.ipynb"))
 
     def test_parse_notebook_with_missing_cell_fields(self, tmp_path):
         """Test parsing notebook where cells are missing cell_type or source"""
@@ -169,17 +164,18 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_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"""
@@ -191,6 +187,7 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_path)
         
-        assert result['cells'] == []
+        assert result == []

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -5,14 +5,17 @@ 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
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
-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 +37,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 +47,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 +57,9 @@ 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')
+        absolute_path = AbsoluteNotebookPath('/notebook.ipynb')
+        errors = get_runtime_errors(app_code, absolute_path)
         
         if expected_error is None:
             assert errors is None
@@ -84,19 +88,21 @@ 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, AbsoluteNotebookPath(app_path))
             assert errors is None
 
+class TestValidateApp:
+    """Test cases for validate_app function"""
+
     @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')
+    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, AbsoluteNotebookPath('/notebook.ipynb'))
         
         assert has_validation_error == expected_has_validation_error
         assert expected_error_message in str(errors)