@@ -401,14 +401,19 @@ def test_plot_palantir_results_custom_args(mock_anndata):
401401
402402# Test with AnnData and all keys available
403403def test_plot_terminal_state_probs_anndata (mock_anndata , mock_cells ):
404+ # Close existing figures to avoid warnings
405+ plt .close ('all' )
404406 fig = plot_terminal_state_probs (mock_anndata , mock_cells )
405407 assert isinstance (fig , plt .Figure )
408+ plt .close (fig )
406409
407410
408411# Test with DataFrame and PResults
409412def test_plot_terminal_state_probs_dataframe (mock_data , mock_presults , mock_cells ):
413+ plt .close ('all' )
410414 fig = plot_terminal_state_probs (mock_data , mock_cells , pr_res = mock_presults )
411415 assert isinstance (fig , plt .Figure )
416+ plt .close (fig )
412417
413418
414419# Test ValueError for missing pr_res in DataFrame input
@@ -616,31 +621,42 @@ def test_plot_gene_trend_heatmaps(mock_anndata):
616621
617622
618623def test_plot_gene_trend_clusters (mock_anndata ):
624+ # Close all figures to start clean
625+ plt .close ('all' )
626+
619627 # Test with AnnData object
620628 fig = plot_gene_trend_clusters (mock_anndata , branch_name = "a" , clusters = "clusters" )
621629 assert isinstance (fig , plt .Figure )
622630
623631 # We expect one subplot per unique non-NaN cluster
624632 unique_clusters = [x for x in mock_anndata .var ["clusters" ].unique () if not pd .isna (x )]
625633 assert len (fig .axes ) == len (unique_clusters )
634+ plt .close (fig )
626635
636+ # Close all figures again before next test
637+ plt .close ('all' )
638+
627639 # Test DataFrame input
628640 trends_df = mock_anndata .varm ["gene_trends_a" ]
629641 clusters_series = mock_anndata .var ["clusters" ]
630642 fig_df = plot_gene_trend_clusters (trends_df , clusters = clusters_series )
631643
632644 assert isinstance (fig_df , plt .Figure )
633645 assert len (fig_df .axes ) == len (unique_clusters )
634-
635- plt .close (fig )
646+
636647 plt .close (fig_df )
648+ plt .close ('all' )
637649
638650
639651def test_gene_score_histogram (mock_anndata ):
652+ # Close all figures before testing
653+ plt .close ('all' )
654+
640655 # Test with minimum required parameters
641656 fig = gene_score_histogram (mock_anndata , "gene_score" )
642657 assert isinstance (fig , plt .Figure )
643658 plt .close (fig )
659+ plt .close ('all' )
644660
645661 # Test with optional parameters
646662 fig = gene_score_histogram (
@@ -652,15 +668,18 @@ def test_gene_score_histogram(mock_anndata):
652668 )
653669 assert isinstance (fig , plt .Figure )
654670 plt .close (fig )
671+ plt .close ('all' )
655672
656- # Test with None quantile
673+ # Test with None quantile - this would previously cause a warning
674+ # about empty legend when quantile is None
657675 fig = gene_score_histogram (
658676 mock_anndata ,
659677 "gene_score" ,
660678 quantile = None ,
661679 )
662680 assert isinstance (fig , plt .Figure )
663681 plt .close (fig )
682+ plt .close ('all' )
664683
665684
666685def test_gene_score_histogram_errors (mock_anndata ):
0 commit comments