|
| 1 | +"""Unit tests for population mapping in network visualization.""" |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +from pypopart.core.alignment import Alignment |
| 6 | +from pypopart.core.sequence import Sequence |
| 7 | +from pypopart.algorithms.mst import MinimumSpanningTree |
| 8 | +from pypopart.visualization.cytoscape_plot import ( |
| 9 | + InteractiveCytoscapePlotter, |
| 10 | + create_cytoscape_network, |
| 11 | +) |
| 12 | + |
| 13 | + |
| 14 | +class TestPopulationMapping: |
| 15 | + """Test cases for mapping sample_ids to populations in network visualization.""" |
| 16 | + |
| 17 | + @pytest.fixture |
| 18 | + def sample_sequences(self): |
| 19 | + """Create sample sequences for testing.""" |
| 20 | + sequences = [ |
| 21 | + Sequence(id='seq1', data='ATCG'), |
| 22 | + Sequence(id='seq2', data='ATCG'), # Same as seq1 |
| 23 | + Sequence(id='seq3', data='ATCG'), # Same as seq1 |
| 24 | + Sequence(id='seq4', data='TTCG'), # Different |
| 25 | + Sequence(id='seq5', data='TTCG'), # Same as seq4 |
| 26 | + ] |
| 27 | + return sequences |
| 28 | + |
| 29 | + @pytest.fixture |
| 30 | + def sample_alignment(self, sample_sequences): |
| 31 | + """Create a sample alignment.""" |
| 32 | + return Alignment(sample_sequences) |
| 33 | + |
| 34 | + @pytest.fixture |
| 35 | + def sample_network(self, sample_alignment): |
| 36 | + """Create a sample network.""" |
| 37 | + algo = MinimumSpanningTree(distance_metric='hamming') |
| 38 | + network = algo.build_network(sample_alignment) |
| 39 | + return network |
| 40 | + |
| 41 | + @pytest.fixture |
| 42 | + def population_mapping(self): |
| 43 | + """Create a population mapping for sequences.""" |
| 44 | + return { |
| 45 | + 'seq1': 'PopA', |
| 46 | + 'seq2': 'PopA', |
| 47 | + 'seq3': 'PopB', |
| 48 | + 'seq4': 'PopB', |
| 49 | + 'seq5': 'PopC', |
| 50 | + } |
| 51 | + |
| 52 | + @pytest.fixture |
| 53 | + def population_colors(self): |
| 54 | + """Create population colors.""" |
| 55 | + return { |
| 56 | + 'PopA': '#FF0000', |
| 57 | + 'PopB': '#00FF00', |
| 58 | + 'PopC': '#0000FF', |
| 59 | + } |
| 60 | + |
| 61 | + def test_create_elements_with_population_mapping( |
| 62 | + self, sample_network, population_mapping, population_colors |
| 63 | + ): |
| 64 | + """Test that create_elements uses population_mapping correctly.""" |
| 65 | + plotter = InteractiveCytoscapePlotter(sample_network) |
| 66 | + |
| 67 | + # Create elements with population mapping |
| 68 | + elements = plotter.create_elements( |
| 69 | + population_colors=population_colors, |
| 70 | + population_mapping=population_mapping, |
| 71 | + ) |
| 72 | + |
| 73 | + # Filter to only node elements |
| 74 | + nodes = [e for e in elements if 'source' not in e.get('data', {})] |
| 75 | + |
| 76 | + # Check that we have nodes |
| 77 | + assert len(nodes) > 0 |
| 78 | + |
| 79 | + # Find nodes with multiple populations (should have pie charts) |
| 80 | + pie_nodes = [n for n in nodes if n['data'].get('has_pie', False)] |
| 81 | + |
| 82 | + # We expect at least one node to have multiple populations |
| 83 | + # (haplotype with seq1, seq2, seq3 has PopA and PopB) |
| 84 | + assert len(pie_nodes) > 0 |
| 85 | + |
| 86 | + # Check that pie nodes have SVG data |
| 87 | + for node in pie_nodes: |
| 88 | + assert 'pie_svg' in node['data'] |
| 89 | + assert node['data']['pie_svg'].startswith('data:image/svg+xml;base64,') |
| 90 | + assert node['data']['color'] == 'transparent' |
| 91 | + |
| 92 | + def test_create_elements_without_population_mapping( |
| 93 | + self, sample_network, population_colors |
| 94 | + ): |
| 95 | + """Test that create_elements works without population_mapping.""" |
| 96 | + plotter = InteractiveCytoscapePlotter(sample_network) |
| 97 | + |
| 98 | + # Create elements without population mapping |
| 99 | + elements = plotter.create_elements( |
| 100 | + population_colors=population_colors, |
| 101 | + population_mapping=None, |
| 102 | + ) |
| 103 | + |
| 104 | + # Should still create elements (just won't have population data) |
| 105 | + assert len(elements) > 0 |
| 106 | + |
| 107 | + # Filter to only node elements |
| 108 | + nodes = [e for e in elements if 'source' not in e.get('data', {})] |
| 109 | + assert len(nodes) > 0 |
| 110 | + |
| 111 | + def test_create_cytoscape_network_with_population_mapping( |
| 112 | + self, sample_network, population_mapping, population_colors |
| 113 | + ): |
| 114 | + """Test the wrapper function with population mapping.""" |
| 115 | + elements, stylesheet = create_cytoscape_network( |
| 116 | + sample_network, |
| 117 | + population_colors=population_colors, |
| 118 | + population_mapping=population_mapping, |
| 119 | + ) |
| 120 | + |
| 121 | + # Check elements were created |
| 122 | + assert len(elements) > 0 |
| 123 | + |
| 124 | + # Check stylesheet was created |
| 125 | + assert len(stylesheet) > 0 |
| 126 | + |
| 127 | + # Filter to only node elements |
| 128 | + nodes = [e for e in elements if 'source' not in e.get('data', {})] |
| 129 | + |
| 130 | + # Check for pie chart nodes |
| 131 | + pie_nodes = [n for n in nodes if n['data'].get('has_pie', False)] |
| 132 | + assert len(pie_nodes) > 0 |
| 133 | + |
| 134 | + def test_single_population_node_color( |
| 135 | + self, sample_network, population_mapping, population_colors |
| 136 | + ): |
| 137 | + """Test that nodes with single population get correct color.""" |
| 138 | + plotter = InteractiveCytoscapePlotter(sample_network) |
| 139 | + |
| 140 | + elements = plotter.create_elements( |
| 141 | + population_colors=population_colors, |
| 142 | + population_mapping=population_mapping, |
| 143 | + ) |
| 144 | + |
| 145 | + # Filter to only node elements |
| 146 | + nodes = [e for e in elements if 'source' not in e.get('data', {})] |
| 147 | + |
| 148 | + # Find nodes with single population |
| 149 | + single_pop_nodes = [ |
| 150 | + n for n in nodes if not n['data'].get('has_pie', False) and not n['data'].get('is_median', False) |
| 151 | + ] |
| 152 | + |
| 153 | + # At least one node should have single population |
| 154 | + # Check they have appropriate colors |
| 155 | + for node in single_pop_nodes: |
| 156 | + node_color = node['data'].get('color') |
| 157 | + # Color should either be from population_colors or default |
| 158 | + assert node_color is not None |
| 159 | + |
| 160 | + def test_pie_chart_data_structure( |
| 161 | + self, sample_network, population_mapping, population_colors |
| 162 | + ): |
| 163 | + """Test that pie chart data has correct structure.""" |
| 164 | + plotter = InteractiveCytoscapePlotter(sample_network) |
| 165 | + |
| 166 | + elements = plotter.create_elements( |
| 167 | + population_colors=population_colors, |
| 168 | + population_mapping=population_mapping, |
| 169 | + ) |
| 170 | + |
| 171 | + # Filter to pie chart nodes |
| 172 | + nodes = [e for e in elements if 'source' not in e.get('data', {})] |
| 173 | + pie_nodes = [n for n in nodes if n['data'].get('has_pie', False)] |
| 174 | + |
| 175 | + for node in pie_nodes: |
| 176 | + # Check pie_data structure |
| 177 | + assert 'pie_data' in node['data'] |
| 178 | + pie_data = node['data']['pie_data'] |
| 179 | + assert isinstance(pie_data, list) |
| 180 | + |
| 181 | + for segment in pie_data: |
| 182 | + assert 'population' in segment |
| 183 | + assert 'value' in segment |
| 184 | + assert 'percent' in segment |
| 185 | + assert 'color' in segment |
| 186 | + # Percent should sum to 100 across all segments |
| 187 | + # Value should be positive |
| 188 | + assert segment['value'] > 0 |
| 189 | + assert 0 < segment['percent'] <= 100 |
| 190 | + |
| 191 | + def test_hover_text_includes_population_data( |
| 192 | + self, sample_network, population_mapping, population_colors |
| 193 | + ): |
| 194 | + """Test that hover text includes population information.""" |
| 195 | + plotter = InteractiveCytoscapePlotter(sample_network) |
| 196 | + |
| 197 | + elements = plotter.create_elements( |
| 198 | + population_colors=population_colors, |
| 199 | + population_mapping=population_mapping, |
| 200 | + ) |
| 201 | + |
| 202 | + # Filter to only non-median node elements |
| 203 | + nodes = [ |
| 204 | + e |
| 205 | + for e in elements |
| 206 | + if 'source' not in e.get('data', {}) and not e['data'].get('is_median', False) |
| 207 | + ] |
| 208 | + |
| 209 | + # Check that hover text exists and contains population info |
| 210 | + for node in nodes: |
| 211 | + # Verify hover text exists |
| 212 | + assert 'hover' in node['data'] |
| 213 | + # Should contain 'Populations:' if there's population data |
| 214 | + if population_mapping and 'Populations:' in node['data']['hover']: |
| 215 | + # Verify population info is in the hover text |
| 216 | + assert node['data']['hover'] is not None |
0 commit comments