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

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -12,7 +12,7 @@ from mito_ai.streamlit_conversion.streamlit_agent_handler import (
     correct_error_in_generation,
     streamlit_handler
 )
-from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
+from mito_ai.path_utils import AbsoluteNotebookPath, AppFileName, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
 
 # Add this line to enable async support
 pytest_plugins = ('pytest_asyncio',)
@@ -102,14 +102,12 @@ class TestCorrectErrorInGeneration:
     @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]:
@@ -117,8 +115,8 @@ class TestCorrectErrorInGeneration:
 
         mock_stream.return_value = mock_async_gen()
 
-        result = await 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
 
@@ -140,8 +138,8 @@ class TestStreamlitHandler:
     @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_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_conversion_success')
+    async def test_streamlit_handler_success(self, mock_log_success, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test successful streamlit handler execution"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": [{"cell_type": "code", "source": ["import pandas"]}]}]
@@ -151,35 +149,30 @@ class TestStreamlitHandler:
         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
+        mock_validator.return_value = []
         
         # Use a relative path that will work cross-platform
-        notebook_path = "notebook.ipynb"
-        result = await streamlit_handler(notebook_path)
+        notebook_path = AbsoluteNotebookPath("absolute/path/to/notebook.ipynb")
+        app_file_name = AppFileName('test-app-file-name.py')
         
-        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, app_file_name)
+        await streamlit_handler(notebook_path, app_file_name)
         
         # Verify calls
         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')")
+        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.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_correct_error, mock_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_validation_retry')
+    async def test_streamlit_handler_max_retries_exceeded(self, mock_log_retry, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
         """Test streamlit handler when max retries are exceeded"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": []}]
@@ -189,16 +182,15 @@ class TestStreamlitHandler:
         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"])
+        # Mock validation (always errors) - validate_app returns List[str]
+        mock_validator.return_value = ["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"), AppFileName('test-app-file-name.py'))
         
-        # Verify that error correction was called 5 times (max retries)
+        # Verify that error correction was called 5 times (once per error, 5 retries)
+        # Each retry processes 1 error, so 5 retries = 5 calls
         assert mock_correct_error.call_count == 5
 
     @pytest.mark.asyncio
@@ -206,8 +198,7 @@ class TestStreamlitHandler:
     @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_generate_code, 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: List[dict] = [{"cells": []}]
@@ -216,38 +207,30 @@ class TestStreamlitHandler:
         # Mock code generation
         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 validation (no errors) - validate_app returns List[str]
+        mock_validator.return_value = []
         
-        # 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("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):
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
 
     @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("notebook.ipynb")
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
 
     @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.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_generate_code, mock_parse):
+    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: List[dict] = [{"cells": []}]
@@ -256,11 +239,8 @@ class TestStreamlitHandler:
         # Mock code generation failure
         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("notebook.ipynb")
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
 
 
 

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,7 +121,8 @@ 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) == 2
@@ -143,8 +137,9 @@ class TestParseJupyterNotebookToExtractRequiredContent:
 
     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,7 +164,8 @@ 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) == 3
         assert result[0]['cell_type'] == 'code'
@@ -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 == []

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -7,10 +7,10 @@ from unittest.mock import patch, MagicMock
 from mito_ai.streamlit_conversion.validate_streamlit_app import (
     get_syntax_error,
     get_runtime_errors,
-    check_for_errors,
     validate_app
 )
 import pytest
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
 class TestGetSyntaxError:
@@ -19,33 +19,33 @@ class TestGetSyntaxError:
     @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"
+                "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"
+                "import streamlit\nst.title('Hello World'",
+                "SyntaxError",
+                "invalid Python code"
         ),
         # Empty streamlit app is valid
         (
-            "",
-            None,
-            "empty code"
+                "",
+                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"""
 
@@ -56,7 +56,9 @@ class TestGetRuntimeErrors:
     ])
     def test_get_runtime_errors(self, app_code, expected_error):
         """Test getting runtime errors"""
-        errors = 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
@@ -85,23 +87,26 @@ df=pd.read_csv('data.csv')
             with open(csv_path, "w") as f:
                 f.write("name,age\nJohn,25\nJane,30")
                
-            errors = 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", [
+    @pytest.mark.parametrize("app_code,expected_has_errors,expected_error_message", [
         ("x=5", False, ""),
         ("1/0", True, "division by zero"),
-        ("print('Hello World'", True, "SyntaxError"),
+        # Syntax errors are caught during runtime
         ("", False, ""),
     ])
-    def test_validate_app(self, app_code, expected_has_validation_error, expected_error_message):
+    def test_validate_app(self, app_code, expected_has_errors, expected_error_message):
         """Test the validate_app function"""
-        has_validation_error, errors = validate_app(app_code, '/app.py')
+        errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
         
-        assert has_validation_error == expected_has_validation_error
-        assert expected_error_message in str(errors)
+        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

@@ -4,85 +4,115 @@
 import pytest
 import os
 import tempfile
-from unittest.mock import patch, AsyncMock, MagicMock
-from mito_ai.streamlit_preview.utils import ensure_app_exists
+from unittest.mock import patch, Mock, AsyncMock, MagicMock
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
-class TestEnsureAppExists:
-    """Test cases for the ensure_app_exists function."""
+class TestStreamlitPreviewHandler:
+    """Test cases for the StreamlitPreviewHandler."""
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
-        "app_exists,streamlit_handler_success,expected_success,expected_error,streamlit_handler_called,streamlit_handler_return",
+        "app_exists,force_recreate,streamlit_handler_called",
         [
-            # Test case 1: App exists, should use existing file
-            (
-                True,  # app_exists
-                True,  # streamlit_handler_success (not relevant)
-                True,  # expected_success
-                "",    # expected_error
-                False, # streamlit_handler_called
-                None,  # streamlit_handler_return (not used)
-            ),
-            # Test case 2: App doesn't exist, streamlit_handler succeeds
-            (
-                False,  # app_exists
-                True,   # streamlit_handler_success
-                True,   # expected_success
-                "",     # expected_error
-                True,   # streamlit_handler_called
-                (True, "/path/to/app.py", "Success"),  # streamlit_handler_return
-            )
+            # 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_uses_existing_file",
-            "app_does_not_exist_generates_new_one_success",
+            "app_exists_no_force_recreate",
+            "app_does_not_exist_generates_new_one",
+            "app_exists_force_recreate",
         ]
     )
-    async def test_ensure_app_exists(
+    async def test_post_handler_app_generation(
         self,
         app_exists,
-        streamlit_handler_success,
-        expected_success,
-        expected_error,
+        force_recreate,
         streamlit_handler_called,
-        streamlit_handler_return,
     ):
-        """Test ensure_app_exists function with various scenarios."""
+        """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)
             
-            # Set up app_path based on whether app exists
-            app_path = os.path.join(temp_dir, "app.py") if app_exists else None
+            # Create notebook file
+            with open(notebook_path, "w") as f:
+                f.write('{"cells": []}')
             
-            # Create app.py file if it should exist
+            # Create app file if it should exist
             if app_exists:
-                assert app_path is not None  # Type assertion for mypy
                 with open(app_path, "w") as f:
                     f.write("import streamlit as st\nst.write('Hello World')")
             
-            # Mock get_app_path to return the appropriate value
-            with patch('mito_ai.streamlit_preview.utils.get_app_path') as mock_get_app_path:
-                mock_get_app_path.return_value = app_path
+            # 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()
                 
-                # Mock streamlit_handler
-                with patch('mito_ai.streamlit_preview.utils.streamlit_handler') as mock_streamlit_handler:
-                    if streamlit_handler_return is not None:
-                        mock_streamlit_handler.return_value = streamlit_handler_return
-                    
-                    success, error_msg = await ensure_app_exists(notebook_path, False, "")
-                    
-                    # Assertions
-                    assert success == expected_success
-                    assert error_msg == expected_error
-                    
-                    # Verify get_app_path was called with the correct directory
-                    mock_get_app_path.assert_called_once_with(temp_dir)
-                    
-                    # Verify streamlit_handler was called or not called as expected
-                    if streamlit_handler_called:
-                        mock_streamlit_handler.assert_called_once_with(notebook_path, "")
-                    else:
-                        mock_streamlit_handler.assert_not_called()
+                # 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