tianxuan/tests/test_pipeline.py

"""Test the full pipeline on multiple datasets × algorithms.

Runs 6 parameterized test cases: 3 datasets × 2 algorithms.
Uses real CSV files from the ``data/`` directory.
"""

from io import StringIO

from django.test import TestCase

from analysis.management.commands.run_pipeline import Command as PipelineCommand
from analysis.models import AnalysisRun
from analysis.session_store import SessionStore


class PipelineMultiDatasetTest(TestCase):
    """Run the pipeline on 3 datasets × 2 algorithms = 6 test cases."""

    DATASETS = [
        ('data/test_flows.csv', 'simple_500'),
        ('data/e2e_test.csv', 'e2e_500'),
        ('data/complex_test.csv', 'complex_5000'),
    ]
    ALGORITHMS = ['hdbscan', 'kmeans']

    def setUp(self):
        """Clean SessionStore before each test."""
        SessionStore().drop_all()

    def tearDown(self):
        """Clean SessionStore after each test."""
        SessionStore().drop_all()

    # ── helpers ──────────────────────────────────────────────────────────

    def _run_pipeline(self, csv_glob, algo):
        """Run the full pipeline via management command.

        Returns
        -------
        tuple[AnalysisRun | None, str]
            The latest ``AnalysisRun`` for *csv_glob* and the full log output.
        """
        out = StringIO()
        cmd = PipelineCommand(stdout=out, stderr=out)
        cmd.handle(csv_glob=csv_glob, entity_col=None, algo=algo, output=None)
        run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-id').first()
        return run, out.getvalue()

    # ── individual test cases: 3 × 2 = 6 ────────────────────────────────

    def test_simple_hdbscan(self):
        """simple_500 dataset + hdbscan."""
        run, log = self._run_pipeline('data/test_flows.csv', 'hdbscan')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')
        self.assertIsNotNone(run.entity_column)
        self.assertGreater(run.entity_count, 0)

    def test_simple_kmeans(self):
        """simple_500 dataset + kmeans."""
        run, log = self._run_pipeline('data/test_flows.csv', 'kmeans')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')
        self.assertGreater(run.entity_count, 0)

    def test_e2e_hdbscan(self):
        """e2e_500 dataset + hdbscan."""
        run, log = self._run_pipeline('data/e2e_test.csv', 'hdbscan')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')

    def test_e2e_kmeans(self):
        """e2e_500 dataset + kmeans."""
        run, log = self._run_pipeline('data/e2e_test.csv', 'kmeans')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')

    def test_complex_hdbscan(self):
        """complex_5000 dataset + hdbscan."""
        run, log = self._run_pipeline('data/complex_test.csv', 'hdbscan')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')

    def test_complex_kmeans(self):
        """complex_5000 dataset + kmeans."""
        run, log = self._run_pipeline('data/complex_test.csv', 'kmeans')
        self.assertEqual(run.status, 'completed',
                         f'Pipeline failed: {run.error_message}')
        self.assertGreater(run.cluster_count, 0, 'No clusters found')

    # ── combined parametrised run ────────────────────────────────────────

    def test_all_datasets_all_algorithms(self):
        """Every (dataset, algorithm) combination succeeds."""
        failures = []
        for csv_glob, label in self.DATASETS:
            for algo in self.ALGORITHMS:
                with self.subTest(dataset=label, algo=algo):
                    run, log = self._run_pipeline(csv_glob, algo)
                    if run.status != 'completed':
                        failures.append(f'{label}/{algo}: {run.error_message}')
                    self.assertEqual(
                        run.status, 'completed',
                        f'{label}/{algo} failed: {run.error_message}',
                    )
                    self.assertGreater(
                        run.cluster_count, 0,
                        f'{label}/{algo}: no clusters',
                    )
        self.assertEqual(len(failures), 0,
                         f'{len(failures)} failures: {failures}')

    # ── edge cases ──────────────────────────────────────────────────────

    def test_invalid_csv_glob(self):
        """Non-existent CSV → pipeline reports failure, does not crash."""
        out = StringIO()
        cmd = PipelineCommand(stdout=out, stderr=out)
        cmd.handle(csv_glob='data/nonexistent.csv',
                   entity_col=None, algo='hdbscan', output=None)
        run = AnalysisRun.objects.filter(
            csv_glob='data/nonexistent.csv',
        ).order_by('-id').first()
        self.assertIsNotNone(run)
        # Should fail gracefully (file not found → no rows → no entity col)
        self.assertEqual(run.status, 'failed')