diff --git a/pystackql/magic_ext/base.py b/pystackql/magic_ext/base.py
index 45e2f3c..b48d770 100644
--- a/pystackql/magic_ext/base.py
+++ b/pystackql/magic_ext/base.py
@@ -7,8 +7,9 @@
"""
from __future__ import print_function
-from IPython.core.magic import Magics
+from IPython.core.magic import Magics, line_cell_magic
from string import Template
+import argparse
class BaseStackqlMagic(Magics):
"""Base Jupyter magic extension enabling running StackQL queries.
@@ -26,6 +27,7 @@ def __init__(self, shell, server_mode):
from ..core import StackQL
super(BaseStackqlMagic, self).__init__(shell)
self.stackql_instance = StackQL(server_mode=server_mode, output='pandas')
+ self.server_mode = server_mode
def get_rendered_query(self, data):
"""Substitute placeholders in a query template with variables from the current namespace.
@@ -52,10 +54,53 @@ def run_query(self, query):
return self.stackql_instance.execute(query)
+ @line_cell_magic
+ def stackql(self, line, cell=None):
+ """A Jupyter magic command to run StackQL queries.
+
+ Can be used as both line and cell magic:
+ - As a line magic: `%stackql QUERY`
+ - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
+
+ :param line: The arguments and/or StackQL query when used as line magic.
+ :param cell: The StackQL query when used as cell magic.
+ :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
+ """
+ is_cell_magic = cell is not None
+
+ if is_cell_magic:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
+ parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
+ args = parser.parse_args(line.split())
+ query_to_run = self.get_rendered_query(cell)
+ else:
+ args = None
+ query_to_run = self.get_rendered_query(line)
+
+ results = self.run_query(query_to_run)
+ self.shell.user_ns['stackql_df'] = results
+
+ if is_cell_magic and args and args.no_display:
+ return None
+ elif is_cell_magic and args and args.csv_download and not args.no_display:
+ # First display the DataFrame
+ import IPython.display
+ IPython.display.display(results)
+ # Then add the download button without displaying the DataFrame again
+ self._display_with_csv_download(results)
+ return results
+ elif is_cell_magic and args and not args.no_display:
+ return results
+ elif not is_cell_magic:
+ return results
+ else:
+ return results
+
def _display_with_csv_download(self, df):
- """Display DataFrame with CSV download link.
+ """Display a CSV download link for the DataFrame without displaying the DataFrame again.
- :param df: The DataFrame to display and make downloadable.
+ :param df: The DataFrame to make downloadable.
"""
import IPython.display
@@ -73,22 +118,7 @@ def _display_with_csv_download(self, df):
# Create download link
download_link = f'data:text/csv;base64,{csv_base64}'
- # Display the DataFrame first
- IPython.display.display(df)
-
- # # Create and display the download button
- # download_html = f'''
- #
- # '''
-
- # Create and display the download button
+ # Only display the download button, not the DataFrame
download_html = f'''
- '''
-
+ '''
IPython.display.display(IPython.display.HTML(download_html))
except Exception as e:
- # If CSV generation fails, just display the DataFrame normally
- IPython.display.display(df)
+ # If CSV generation fails, just print an error message without displaying anything
print(f"Error generating CSV download: {e}")
\ No newline at end of file
diff --git a/pystackql/magic_ext/local.py b/pystackql/magic_ext/local.py
index e88479b..1830249 100644
--- a/pystackql/magic_ext/local.py
+++ b/pystackql/magic_ext/local.py
@@ -7,9 +7,8 @@
using a local StackQL binary.
"""
-from IPython.core.magic import (magics_class, line_cell_magic)
+from IPython.core.magic import magics_class
from .base import BaseStackqlMagic
-import argparse
@magics_class
class StackqlMagic(BaseStackqlMagic):
@@ -22,43 +21,6 @@ def __init__(self, shell):
"""
super().__init__(shell, server_mode=False)
- @line_cell_magic
- def stackql(self, line, cell=None):
- """A Jupyter magic command to run StackQL queries.
-
- Can be used as both line and cell magic:
- - As a line magic: `%stackql QUERY`
- - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
-
- :param line: The arguments and/or StackQL query when used as line magic.
- :param cell: The StackQL query when used as cell magic.
- :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
- """
- is_cell_magic = cell is not None
-
- if is_cell_magic:
- parser = argparse.ArgumentParser()
- parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
- parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
- args = parser.parse_args(line.split())
- query_to_run = self.get_rendered_query(cell)
- else:
- args = None
- query_to_run = self.get_rendered_query(line)
-
- results = self.run_query(query_to_run)
- self.shell.user_ns['stackql_df'] = results
-
- if is_cell_magic and args and args.no_display:
- return None
- elif is_cell_magic and args and args.csv_download and not args.no_display:
- self._display_with_csv_download(results)
- return results
- elif is_cell_magic and args and not args.no_display:
- return results
- elif not is_cell_magic:
- return results
-
def load_ipython_extension(ipython):
"""Load the non-server magic in IPython.
diff --git a/pystackql/magic_ext/server.py b/pystackql/magic_ext/server.py
index 70d561c..2c6d8f1 100644
--- a/pystackql/magic_ext/server.py
+++ b/pystackql/magic_ext/server.py
@@ -7,9 +7,8 @@
using a StackQL server connection.
"""
-from IPython.core.magic import (magics_class, line_cell_magic)
+from IPython.core.magic import magics_class
from .base import BaseStackqlMagic
-import argparse
@magics_class
class StackqlServerMagic(BaseStackqlMagic):
@@ -22,45 +21,6 @@ def __init__(self, shell):
"""
super().__init__(shell, server_mode=True)
- @line_cell_magic
- def stackql(self, line, cell=None):
- """A Jupyter magic command to run StackQL queries.
-
- Can be used as both line and cell magic:
- - As a line magic: `%stackql QUERY`
- - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
-
- :param line: The arguments and/or StackQL query when used as line magic.
- :param cell: The StackQL query when used as cell magic.
- :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
- """
- is_cell_magic = cell is not None
-
- if is_cell_magic:
- parser = argparse.ArgumentParser()
- parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
- parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
- args = parser.parse_args(line.split())
- query_to_run = self.get_rendered_query(cell)
- else:
- args = None
- query_to_run = self.get_rendered_query(line)
-
- results = self.run_query(query_to_run)
- self.shell.user_ns['stackql_df'] = results
-
- if is_cell_magic and args and args.no_display:
- return None
- elif is_cell_magic and args and args.csv_download and not args.no_display:
- self._display_with_csv_download(results)
- return results
- elif is_cell_magic and args and not args.no_display:
- return results
- elif not is_cell_magic:
- return results
- else:
- return results
-
def load_ipython_extension(ipython):
"""Load the server magic in IPython."""
# Create an instance of the magic class and register it
diff --git a/tests/test_magic.py b/tests/test_magic.py
index aaaac49..edf1bdf 100644
--- a/tests/test_magic.py
+++ b/tests/test_magic.py
@@ -8,214 +8,27 @@
import os
import sys
-import re
import pytest
-import pandas as pd
-from unittest.mock import MagicMock
# Add the parent directory to the path so we can import from pystackql
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-# Add the current directory to the path so we can import test_constants
-sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+# Import the base test class
+from tests.test_magic_base import BaseStackQLMagicTest
# Import directly from the original modules - this is what notebooks would do
from pystackql import magic
from pystackql import StackqlMagic
-from tests.test_constants import (
- LITERAL_INT_QUERY,
- REGISTRY_PULL_HOMEBREW_QUERY,
- registry_pull_resp_pattern,
- print_test_result
-)
+from tests.test_constants import print_test_result
-class TestStackQLMagic:
+class TestStackQLMagic(BaseStackQLMagicTest):
"""Tests for the non-server mode magic extension."""
- @pytest.fixture(autouse=True)
- def setup_method(self, mock_interactive_shell):
- """Set up the test environment."""
- self.shell = mock_interactive_shell
-
- # Load the magic extension
- magic.load_ipython_extension(self.shell)
-
- # Create the magic instance
- self.stackql_magic = StackqlMagic(shell=self.shell)
-
- # Set up test data
- self.query = LITERAL_INT_QUERY
- self.expected_result = pd.DataFrame({"literal_int_value": [1]})
- self.statement = REGISTRY_PULL_HOMEBREW_QUERY
-
- def test_line_magic_query(self):
- """Test line magic with a query."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query
- result = self.stackql_magic.stackql(line=self.query, cell=None)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Line magic query test",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- False, True)
-
- def test_cell_magic_query(self):
- """Test cell magic with a query."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query
- result = self.stackql_magic.stackql(line="", cell=self.query)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Cell magic query test",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- False, True)
-
- def test_cell_magic_query_no_display(self):
- """Test cell magic with a query and --no-display option."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query and --no-display option
- result = self.stackql_magic.stackql(line="--no-display", cell=self.query)
-
- # Validate the outcome
- assert result is None, "Result should be None with --no-display option"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Cell magic query test (with --no-display)",
- result is None and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- False, True)
-
- def test_cell_magic_query_csv_download(self):
- """Test cell magic with CSV download functionality."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Mock the _display_with_csv_download method to verify it's called
- self.stackql_magic._display_with_csv_download = MagicMock()
-
- # Execute the magic with --csv-download option
- result = self.stackql_magic.stackql(line="--csv-download", cell=self.query)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- # Verify that _display_with_csv_download was called
- self.stackql_magic._display_with_csv_download.assert_called_once_with(self.expected_result)
-
- print_test_result("Cell magic query test with CSV download",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.stackql_magic._display_with_csv_download.called,
- False, True)
-
- def test_cell_magic_query_csv_download_with_no_display(self):
- """Test that --no-display takes precedence over --csv-download."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Mock the _display_with_csv_download method to verify it's not called
- self.stackql_magic._display_with_csv_download = MagicMock()
-
- # Execute the magic with both --csv-download and --no-display options
- result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
-
- # Validate the outcome
- assert result is None, "Result should be None with --no-display option"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- # Verify that _display_with_csv_download was NOT called
- self.stackql_magic._display_with_csv_download.assert_not_called()
-
- print_test_result("Cell magic query test with CSV download and no-display",
- result is None and
- 'stackql_df' in self.shell.user_ns and
- not self.stackql_magic._display_with_csv_download.called,
- False, True)
-
- def test_display_with_csv_download_method(self):
- """Test the _display_with_csv_download method directly."""
- import base64
- from unittest.mock import patch
-
- # Create a test DataFrame
- test_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
-
- # Mock IPython display functionality
- with patch('IPython.display.display') as mock_display, \
- patch('IPython.display.HTML') as mock_html:
-
- # Call the method
- self.stackql_magic._display_with_csv_download(test_df)
-
- # Verify display was called twice (once for DataFrame, once for HTML)
- assert mock_display.call_count == 2, "Display should be called twice"
-
- # Verify HTML was called once
- mock_html.assert_called_once()
-
- # Check that the HTML call contains download link
- html_call_args = mock_html.call_args[0][0]
- assert 'download="stackql_results.csv"' in html_call_args
- assert 'data:text/csv;base64,' in html_call_args
-
- print_test_result("_display_with_csv_download method test",
- mock_display.call_count == 2 and mock_html.called,
- False, True)
-
- def test_display_with_csv_download_error_handling(self):
- """Test error handling in _display_with_csv_download method."""
- from unittest.mock import patch
-
- # Create a mock DataFrame that will raise an exception during to_csv()
- mock_df = MagicMock()
- mock_df.to_csv.side_effect = Exception("Test CSV error")
-
- # Mock IPython display functionality
- with patch('IPython.display.display') as mock_display, \
- patch('IPython.display.HTML') as mock_html, \
- patch('builtins.print') as mock_print:
-
- # Call the method with the problematic DataFrame
- self.stackql_magic._display_with_csv_download(mock_df)
-
- # Verify display was called once (for DataFrame only, not for HTML)
- mock_display.assert_called_once_with(mock_df)
-
- # Verify HTML was not called due to error
- mock_html.assert_not_called()
-
- # Verify error message was printed
- mock_print.assert_called_once()
- error_message = mock_print.call_args[0][0]
- assert "Error generating CSV download:" in error_message
-
- print_test_result("_display_with_csv_download error handling test",
- mock_display.called and not mock_html.called and mock_print.called,
- False, True)
+ # Set the class attributes for the base test class
+ magic_module = magic
+ magic_class = StackqlMagic
+ is_server_mode = False
def test_magic_extension_loading(mock_interactive_shell):
"""Test that non-server magic extension can be loaded."""
diff --git a/tests/test_magic_base.py b/tests/test_magic_base.py
new file mode 100644
index 0000000..e5e7b71
--- /dev/null
+++ b/tests/test_magic_base.py
@@ -0,0 +1,221 @@
+# tests/test_magic_base.py
+
+"""
+Base test class for Jupyter magic extensions for PyStackQL.
+
+This module provides a base test class for testing both local and server mode
+magic extensions.
+"""
+
+import os
+import sys
+import re
+import pytest
+import pandas as pd
+from unittest.mock import MagicMock, patch
+
+# Add the parent directory to the path so we can import from pystackql
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# Add the current directory to the path so we can import test_constants
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+
+from tests.test_constants import (
+ LITERAL_INT_QUERY,
+ REGISTRY_PULL_HOMEBREW_QUERY,
+ registry_pull_resp_pattern,
+ print_test_result
+)
+
+class BaseStackQLMagicTest:
+ """Base class for testing StackQL magic extensions."""
+
+ # Each derived class should define:
+ # - magic_module: the module to import
+ # - magic_class: the class to use
+ # - is_server_mode: True for server mode tests, False for local mode tests
+ magic_module = None
+ magic_class = None
+ is_server_mode = None
+
+ @pytest.fixture(autouse=True)
+ def setup_method(self, mock_interactive_shell):
+ """Set up the test environment."""
+ self.shell = mock_interactive_shell
+
+ # Load the magic extension
+ self.magic_module.load_ipython_extension(self.shell)
+
+ # Create the magic instance
+ self.stackql_magic = self.magic_class(shell=self.shell)
+
+ # Set up test data
+ self.query = LITERAL_INT_QUERY
+ self.expected_result = pd.DataFrame({"literal_int_value": [1]})
+ self.statement = REGISTRY_PULL_HOMEBREW_QUERY
+
+ def test_line_magic_query(self):
+ """Test line magic with a query."""
+ # Mock the run_query method to return a known DataFrame
+ self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+ # Execute the magic with our query
+ result = self.stackql_magic.stackql(line=self.query, cell=None)
+
+ # Validate the outcome
+ assert result.equals(self.expected_result), "Result should match expected DataFrame"
+ assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
+ assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+ print_test_result(f"Line magic query test{' (server mode)' if self.is_server_mode else ''}",
+ result.equals(self.expected_result) and
+ 'stackql_df' in self.shell.user_ns and
+ self.shell.user_ns['stackql_df'].equals(self.expected_result),
+ self.is_server_mode, True)
+
+ def test_cell_magic_query(self):
+ """Test cell magic with a query."""
+ # Mock the run_query method to return a known DataFrame
+ self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+ # Execute the magic with our query
+ result = self.stackql_magic.stackql(line="", cell=self.query)
+
+ # Validate the outcome
+ assert result.equals(self.expected_result), "Result should match expected DataFrame"
+ assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
+ assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+ print_test_result(f"Cell magic query test{' (server mode)' if self.is_server_mode else ''}",
+ result.equals(self.expected_result) and
+ 'stackql_df' in self.shell.user_ns and
+ self.shell.user_ns['stackql_df'].equals(self.expected_result),
+ self.is_server_mode, True)
+
+ def test_cell_magic_query_no_display(self):
+ """Test cell magic with a query and --no-display option."""
+ # Mock the run_query method to return a known DataFrame
+ self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+ # Execute the magic with our query and --no-display option
+ result = self.stackql_magic.stackql(line="--no-display", cell=self.query)
+
+ # Validate the outcome
+ assert result is None, "Result should be None with --no-display option"
+ assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
+ assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+ print_test_result(f"Cell magic query test with --no-display{' (server mode)' if self.is_server_mode else ''}",
+ result is None and
+ 'stackql_df' in self.shell.user_ns and
+ self.shell.user_ns['stackql_df'].equals(self.expected_result),
+ self.is_server_mode, True)
+
+ def test_cell_magic_query_csv_download(self):
+ """Test cell magic with CSV download functionality."""
+ # Mock the run_query method to return a known DataFrame
+ self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+ # Mock the _display_with_csv_download method to verify it's called
+ self.stackql_magic._display_with_csv_download = MagicMock()
+
+ # Execute the magic with --csv-download option
+ result = self.stackql_magic.stackql(line="--csv-download", cell=self.query)
+
+ # Validate the outcome
+ assert result.equals(self.expected_result), "Result should match expected DataFrame"
+ assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
+ assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+ # Verify that _display_with_csv_download was called
+ self.stackql_magic._display_with_csv_download.assert_called_once_with(self.expected_result)
+
+ print_test_result(f"Cell magic query test with CSV download{' (server mode)' if self.is_server_mode else ''}",
+ result.equals(self.expected_result) and
+ 'stackql_df' in self.shell.user_ns and
+ self.stackql_magic._display_with_csv_download.called,
+ self.is_server_mode, True)
+
+ def test_cell_magic_query_csv_download_with_no_display(self):
+ """Test that --no-display takes precedence over --csv-download."""
+ # Mock the run_query method to return a known DataFrame
+ self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+ # Mock the _display_with_csv_download method to verify it's not called
+ self.stackql_magic._display_with_csv_download = MagicMock()
+
+ # Execute the magic with both --csv-download and --no-display options
+ result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
+
+ # Validate the outcome
+ assert result is None, "Result should be None with --no-display option"
+ assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
+ assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+ # Verify that _display_with_csv_download was NOT called
+ self.stackql_magic._display_with_csv_download.assert_not_called()
+
+ print_test_result(f"Cell magic query test with CSV download and no-display{' (server mode)' if self.is_server_mode else ''}",
+ result is None and
+ 'stackql_df' in self.shell.user_ns and
+ not self.stackql_magic._display_with_csv_download.called,
+ self.is_server_mode, True)
+
+ def test_display_with_csv_download_method(self):
+ """Test the _display_with_csv_download method directly."""
+ import base64
+
+ # Create a test DataFrame
+ test_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
+
+ # Mock IPython display functionality
+ with patch('IPython.display.display') as mock_display, \
+ patch('IPython.display.HTML') as mock_html:
+
+ # Call the method
+ self.stackql_magic._display_with_csv_download(test_df)
+
+ # Verify display was called once (only for HTML, not for DataFrame)
+ assert mock_display.call_count == 1, "Display should be called once"
+
+ # Verify HTML was called once
+ mock_html.assert_called_once()
+
+ # Check that the HTML call contains download link
+ html_call_args = mock_html.call_args[0][0]
+ assert 'download="stackql_results.csv"' in html_call_args
+ assert 'data:text/csv;base64,' in html_call_args
+
+ print_test_result(f"_display_with_csv_download method test{' (server mode)' if self.is_server_mode else ''}",
+ mock_display.call_count == 1 and mock_html.called,
+ self.is_server_mode, True)
+
+ def test_display_with_csv_download_error_handling(self):
+ """Test error handling in _display_with_csv_download method."""
+
+ # Create a mock DataFrame that will raise an exception during to_csv()
+ mock_df = MagicMock()
+ mock_df.to_csv.side_effect = Exception("Test CSV error")
+
+ # Mock IPython display functionality
+ with patch('IPython.display.display') as mock_display, \
+ patch('IPython.display.HTML') as mock_html, \
+ patch('builtins.print') as mock_print:
+
+ # Call the method with the problematic DataFrame
+ self.stackql_magic._display_with_csv_download(mock_df)
+
+ # Verify display was not called in the error case
+ mock_display.assert_not_called()
+
+ # Verify HTML was not called in the error case
+ mock_html.assert_not_called()
+
+ # Verify error message was printed
+ mock_print.assert_called_once()
+ error_message = mock_print.call_args[0][0]
+ assert "Error generating CSV download:" in error_message
+
+ print_test_result(f"_display_with_csv_download error handling test{' (server mode)' if self.is_server_mode else ''}",
+ not mock_display.called and not mock_html.called and mock_print.called,
+ self.is_server_mode, True)
\ No newline at end of file
diff --git a/tests/test_server_magic.py b/tests/test_server_magic.py
index 0296f37..d39084d 100644
--- a/tests/test_server_magic.py
+++ b/tests/test_server_magic.py
@@ -1,3 +1,5 @@
+# tests/test_server_magic.py
+
"""
Server-mode magic extension tests for PyStackQL.
@@ -6,214 +8,27 @@
import os
import sys
-import re
import pytest
-import pandas as pd
-from unittest.mock import MagicMock
# Add the parent directory to the path so we can import from pystackql
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-# Add the current directory to the path so we can import test_constants
-sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+# Import the base test class
+from tests.test_magic_base import BaseStackQLMagicTest
# Import directly from the original modules - this is what notebooks would do
from pystackql import magics
from pystackql import StackqlServerMagic
-from tests.test_constants import (
- LITERAL_INT_QUERY,
- REGISTRY_PULL_HOMEBREW_QUERY,
- registry_pull_resp_pattern,
- print_test_result
-)
+from tests.test_constants import print_test_result
-class TestStackQLServerMagic:
+class TestStackQLServerMagic(BaseStackQLMagicTest):
"""Tests for the server mode magic extension."""
- @pytest.fixture(autouse=True)
- def setup_method(self, mock_interactive_shell):
- """Set up the test environment."""
- self.shell = mock_interactive_shell
-
- # Load the magic extension
- magics.load_ipython_extension(self.shell)
-
- # Create the magic instance
- self.stackql_magic = StackqlServerMagic(shell=self.shell)
-
- # Set up test data
- self.query = LITERAL_INT_QUERY
- self.expected_result = pd.DataFrame({"literal_int_value": [1]})
- self.statement = REGISTRY_PULL_HOMEBREW_QUERY
-
- def test_line_magic_query(self):
- """Test line magic with a query in server mode."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query
- result = self.stackql_magic.stackql(line=self.query, cell=None)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Line magic query test (server mode)",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- True, True)
-
- def test_cell_magic_query(self):
- """Test cell magic with a query in server mode."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query
- result = self.stackql_magic.stackql(line="", cell=self.query)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Cell magic query test (server mode)",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- True, True)
-
- def test_cell_magic_query_no_display(self):
- """Test cell magic with a query and --no-display option in server mode."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Execute the magic with our query and --no-display option
- result = self.stackql_magic.stackql(line="--no-display", cell=self.query)
-
- # Validate the outcome
- assert result is None, "Result should be None with --no-display option"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- print_test_result("Cell magic query test with --no-display (server mode)",
- result is None and
- 'stackql_df' in self.shell.user_ns and
- self.shell.user_ns['stackql_df'].equals(self.expected_result),
- True, True)
-
- def test_cell_magic_query_csv_download(self):
- """Test cell magic with CSV download functionality in server mode."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Mock the _display_with_csv_download method to verify it's called
- self.stackql_magic._display_with_csv_download = MagicMock()
-
- # Execute the magic with --csv-download option
- result = self.stackql_magic.stackql(line="--csv-download", cell=self.query)
-
- # Validate the outcome
- assert result.equals(self.expected_result), "Result should match expected DataFrame"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- # Verify that _display_with_csv_download was called
- self.stackql_magic._display_with_csv_download.assert_called_once_with(self.expected_result)
-
- print_test_result("Cell magic query test with CSV download (server mode)",
- result.equals(self.expected_result) and
- 'stackql_df' in self.shell.user_ns and
- self.stackql_magic._display_with_csv_download.called,
- True, True)
-
- def test_cell_magic_query_csv_download_with_no_display(self):
- """Test that --no-display takes precedence over --csv-download in server mode."""
- # Mock the run_query method to return a known DataFrame
- self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
-
- # Mock the _display_with_csv_download method to verify it's not called
- self.stackql_magic._display_with_csv_download = MagicMock()
-
- # Execute the magic with both --csv-download and --no-display options
- result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
-
- # Validate the outcome
- assert result is None, "Result should be None with --no-display option"
- assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
- assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
-
- # Verify that _display_with_csv_download was NOT called
- self.stackql_magic._display_with_csv_download.assert_not_called()
-
- print_test_result("Cell magic query test with CSV download and no-display (server mode)",
- result is None and
- 'stackql_df' in self.shell.user_ns and
- not self.stackql_magic._display_with_csv_download.called,
- True, True)
-
- def test_display_with_csv_download_method(self):
- """Test the _display_with_csv_download method directly in server mode."""
- import base64
- from unittest.mock import patch
-
- # Create a test DataFrame
- test_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
-
- # Mock IPython display functionality
- with patch('IPython.display.display') as mock_display, \
- patch('IPython.display.HTML') as mock_html:
-
- # Call the method
- self.stackql_magic._display_with_csv_download(test_df)
-
- # Verify display was called twice (once for DataFrame, once for HTML)
- assert mock_display.call_count == 2, "Display should be called twice"
-
- # Verify HTML was called once
- mock_html.assert_called_once()
-
- # Check that the HTML call contains download link
- html_call_args = mock_html.call_args[0][0]
- assert 'download="stackql_results.csv"' in html_call_args
- assert 'data:text/csv;base64,' in html_call_args
-
- print_test_result("_display_with_csv_download method test (server mode)",
- mock_display.call_count == 2 and mock_html.called,
- True, True)
-
- def test_display_with_csv_download_error_handling(self):
- """Test error handling in _display_with_csv_download method in server mode."""
- from unittest.mock import patch
-
- # Create a mock DataFrame that will raise an exception during to_csv()
- mock_df = MagicMock()
- mock_df.to_csv.side_effect = Exception("Test CSV error")
-
- # Mock IPython display functionality
- with patch('IPython.display.display') as mock_display, \
- patch('IPython.display.HTML') as mock_html, \
- patch('builtins.print') as mock_print:
-
- # Call the method with the problematic DataFrame
- self.stackql_magic._display_with_csv_download(mock_df)
-
- # Verify display was called once (for DataFrame only, not for HTML)
- mock_display.assert_called_once_with(mock_df)
-
- # Verify HTML was not called due to error
- mock_html.assert_not_called()
-
- # Verify error message was printed
- mock_print.assert_called_once()
- error_message = mock_print.call_args[0][0]
- assert "Error generating CSV download:" in error_message
-
- print_test_result("_display_with_csv_download error handling test (server mode)",
- mock_display.called and not mock_html.called and mock_print.called,
- True, True)
+ # Set the class attributes for the base test class
+ magic_module = magics
+ magic_class = StackqlServerMagic
+ is_server_mode = True
def test_server_magic_extension_loading(mock_interactive_shell):
"""Test that server magic extension can be loaded."""