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