@@ -104,6 +104,117 @@ def test_cell_magic_query_no_display(self):
104104 self .shell .user_ns ['stackql_df' ].equals (self .expected_result ),
105105 True , True )
106106
107+ def test_cell_magic_query_csv_download (self ):
108+ """Test cell magic with CSV download functionality in server mode."""
109+ # Mock the run_query method to return a known DataFrame
110+ self .stackql_magic .run_query = MagicMock (return_value = self .expected_result )
111+
112+ # Mock the _display_with_csv_download method to verify it's called
113+ self .stackql_magic ._display_with_csv_download = MagicMock ()
114+
115+ # Execute the magic with --csv-download option
116+ result = self .stackql_magic .stackql (line = "--csv-download" , cell = self .query )
117+
118+ # Validate the outcome
119+ assert result .equals (self .expected_result ), "Result should match expected DataFrame"
120+ assert 'stackql_df' in self .shell .user_ns , "stackql_df should be in user namespace"
121+ assert self .shell .user_ns ['stackql_df' ].equals (self .expected_result ), "stackql_df should match expected DataFrame"
122+
123+ # Verify that _display_with_csv_download was called
124+ self .stackql_magic ._display_with_csv_download .assert_called_once_with (self .expected_result )
125+
126+ print_test_result ("Cell magic query test with CSV download (server mode)" ,
127+ result .equals (self .expected_result ) and
128+ 'stackql_df' in self .shell .user_ns and
129+ self .stackql_magic ._display_with_csv_download .called ,
130+ True , True )
131+
132+ def test_cell_magic_query_csv_download_with_no_display (self ):
133+ """Test that --no-display takes precedence over --csv-download in server mode."""
134+ # Mock the run_query method to return a known DataFrame
135+ self .stackql_magic .run_query = MagicMock (return_value = self .expected_result )
136+
137+ # Mock the _display_with_csv_download method to verify it's not called
138+ self .stackql_magic ._display_with_csv_download = MagicMock ()
139+
140+ # Execute the magic with both --csv-download and --no-display options
141+ result = self .stackql_magic .stackql (line = "--csv-download --no-display" , cell = self .query )
142+
143+ # Validate the outcome
144+ assert result is None , "Result should be None with --no-display option"
145+ assert 'stackql_df' in self .shell .user_ns , "stackql_df should still be in user namespace"
146+ assert self .shell .user_ns ['stackql_df' ].equals (self .expected_result ), "stackql_df should match expected DataFrame"
147+
148+ # Verify that _display_with_csv_download was NOT called
149+ self .stackql_magic ._display_with_csv_download .assert_not_called ()
150+
151+ print_test_result ("Cell magic query test with CSV download and no-display (server mode)" ,
152+ result is None and
153+ 'stackql_df' in self .shell .user_ns and
154+ not self .stackql_magic ._display_with_csv_download .called ,
155+ True , True )
156+
157+ def test_display_with_csv_download_method (self ):
158+ """Test the _display_with_csv_download method directly in server mode."""
159+ import base64
160+ from unittest .mock import patch
161+
162+ # Create a test DataFrame
163+ test_df = pd .DataFrame ({"col1" : [1 , 2 ], "col2" : ["a" , "b" ]})
164+
165+ # Mock IPython display functionality
166+ with patch ('IPython.display.display' ) as mock_display , \
167+ patch ('IPython.display.HTML' ) as mock_html :
168+
169+ # Call the method
170+ self .stackql_magic ._display_with_csv_download (test_df )
171+
172+ # Verify display was called twice (once for DataFrame, once for HTML)
173+ assert mock_display .call_count == 2 , "Display should be called twice"
174+
175+ # Verify HTML was called once
176+ mock_html .assert_called_once ()
177+
178+ # Check that the HTML call contains download link
179+ html_call_args = mock_html .call_args [0 ][0 ]
180+ assert 'download="stackql_results.csv"' in html_call_args
181+ assert 'data:text/csv;base64,' in html_call_args
182+
183+ print_test_result ("_display_with_csv_download method test (server mode)" ,
184+ mock_display .call_count == 2 and mock_html .called ,
185+ True , True )
186+
187+ def test_display_with_csv_download_error_handling (self ):
188+ """Test error handling in _display_with_csv_download method in server mode."""
189+ from unittest .mock import patch
190+
191+ # Create a mock DataFrame that will raise an exception during to_csv()
192+ mock_df = MagicMock ()
193+ mock_df .to_csv .side_effect = Exception ("Test CSV error" )
194+
195+ # Mock IPython display functionality
196+ with patch ('IPython.display.display' ) as mock_display , \
197+ patch ('IPython.display.HTML' ) as mock_html , \
198+ patch ('builtins.print' ) as mock_print :
199+
200+ # Call the method with the problematic DataFrame
201+ self .stackql_magic ._display_with_csv_download (mock_df )
202+
203+ # Verify display was called once (for DataFrame only, not for HTML)
204+ mock_display .assert_called_once_with (mock_df )
205+
206+ # Verify HTML was not called due to error
207+ mock_html .assert_not_called ()
208+
209+ # Verify error message was printed
210+ mock_print .assert_called_once ()
211+ error_message = mock_print .call_args [0 ][0 ]
212+ assert "Error generating CSV download:" in error_message
213+
214+ print_test_result ("_display_with_csv_download error handling test (server mode)" ,
215+ mock_display .called and not mock_html .called and mock_print .called ,
216+ True , True )
217+
107218def test_server_magic_extension_loading (mock_interactive_shell ):
108219 """Test that server magic extension can be loaded."""
109220 # Test loading server magic
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?v=4&size=40){kind=avatar}
0 commit comments