code/OpenTSLM/test/ecg_qa_cot_test.py

#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT

"""
Test script for the ECG-QA CoT loader and dataset.
"""

import unittest
from opentslm.time_series_datasets.ecg_qa.ECGQACoTQADataset import ECGQACoTQADataset
from opentslm.logger import get_logger, set_global_verbose

def pretty_print_label_distribution(dataset, name):
    """Pretty print label distribution for a dataset."""
    from opentslm.time_series_datasets.ecg_qa.ecgqa_cot_loader import get_label_distribution
    label_dist = get_label_distribution(dataset)
    total = len(dataset)
    print(f"\n{name} dataset:")
    print(f"  Total samples: {total}")
    print(f"  Label distribution:")
    for label, count in sorted(label_dist.items()):
        print(f"    {label:15s}: {count:5d} ({count/total*100:5.1f}%)")

class TestECGQACotLoader(unittest.TestCase):
    """
    Unit tests for the ECG-QA CoT loader functions.
    """
    def setUp(self):
        # Set up global logger with verbose mode for detailed output
        set_global_verbose(True)
        self.logger = get_logger()
        
        from opentslm.time_series_datasets.ecg_qa.ecgqa_cot_loader import load_ecg_qa_cot_splits
        self.load_ecg_qa_cot_splits = load_ecg_qa_cot_splits
        
        self.logger.loading("Loading ECG-QA CoT dataset splits...")
        self.train, self.val, self.test = self.load_ecg_qa_cot_splits()
        self.logger.success(f"Dataset loaded successfully: Train={len(self.train)}, Val={len(self.val)}, Test={len(self.test)}")

    def test_dataset_sizes(self):
        """Test that the datasets are non-empty and splits are correct."""
        self.logger.info("Testing dataset sizes...")
        self.assertGreater(len(self.train), 0)
        self.assertGreater(len(self.val), 0)
        self.assertGreater(len(self.test), 0)
        self.logger.success("Dataset size tests passed")

    def test_label_distribution(self):
        """Test that label distributions are reasonable across splits."""
        self.logger.info("Testing label distributions...")
        pretty_print_label_distribution(self.train, "Train")
        pretty_print_label_distribution(self.val, "Validation")
        pretty_print_label_distribution(self.test, "Test")
        
        from opentslm.time_series_datasets.ecg_qa.ecgqa_cot_loader import get_label_distribution
        train_dist = get_label_distribution(self.train)
        val_dist = get_label_distribution(self.val)
        test_dist = get_label_distribution(self.test)
        
        # Check that all splits have reasonable distributions
        # Note: Not all labels need to appear in all splits due to random sampling
        self.assertGreater(len(train_dist), 0)
        self.assertGreater(len(val_dist), 0)
        self.assertGreater(len(test_dist), 0)
        
        # Check that major labels appear in all splits
        major_labels = ['yes', 'no', 'none']
        for label in major_labels:
            if label in train_dist:
                self.assertIn(label, val_dist, f"Major label '{label}' missing from validation set")
                self.assertIn(label, test_dist, f"Major label '{label}' missing from test set")
        
        self.logger.success("Label distribution tests passed")

    def test_sample_keys(self):
        """Test that a sample contains all required keys."""
        self.logger.info("Testing sample keys...")
        sample = self.train[0]
        required_keys = {"question", "question_type", "template_id", "ecg_id", "ecg_paths", "clinical_contexts", "rationale"}
        self.assertTrue(required_keys.issubset(sample.keys()))
        
        # Test that required fields are not None
        for key in required_keys:
            self.assertIsNotNone(sample[key], f"Required field '{key}' is None")
        
        self.logger.success("Sample keys test passed")

    def test_ecg_data_content(self):
        """Test that the ECG data is present and valid."""
        self.logger.info("Testing ECG data content...")
        sample = self.train[0]
        
        # Check ECG IDs
        self.assertIsInstance(sample["ecg_id"], list)
        self.assertGreater(len(sample["ecg_id"]), 0)
        for ecg_id in sample["ecg_id"]:
            self.assertIsInstance(ecg_id, int)
        
        # Check ECG paths
        self.assertIsInstance(sample["ecg_paths"], list)
        self.assertGreater(len(sample["ecg_paths"]), 0)
        for path in sample["ecg_paths"]:
            self.assertIsInstance(path, str)
            self.assertTrue(path.endswith('.dat'))
        
        # Check clinical contexts
        self.assertIsInstance(sample["clinical_contexts"], list)
        self.assertGreater(len(sample["clinical_contexts"]), 0)
        for context in sample["clinical_contexts"]:
            self.assertIsInstance(context, str)
            self.assertGreater(len(context), 0)
        
        self.logger.success("ECG data content tests passed")

    def test_question_content(self):
        """Test that the question data is valid."""
        self.logger.info("Testing question content...")
        sample = self.train[0]
        
        # Check question
        self.assertIsInstance(sample["question"], str)
        self.assertGreater(len(sample["question"]), 0)
        
        # Check question type
        self.assertIsInstance(sample["question_type"], str)
        self.assertGreater(len(sample["question_type"]), 0)
        
        # Check template ID
        self.assertIsInstance(sample["template_id"], int)
        
        self.logger.success("Question content tests passed")

    def test_rationale_content(self):
        """Test that the rationale is a string and non-empty."""
        self.logger.info("Testing rationale content...")
        sample = self.train[0]
        self.assertIsInstance(sample["rationale"], str)
        self.assertGreater(len(sample["rationale"]), 0)
        self.logger.success(f"Rationale test passed: length={len(sample['rationale'])}")

    def test_cot_fields(self):
        """Test that CoT-specific fields are present."""
        self.logger.info("Testing CoT-specific fields...")
        sample = self.train[0]
        
        # Check CoT fields that actually exist in the implementation
        cot_fields = ["rationale", "template_id", "question_type"]
        for field in cot_fields:
            self.assertIn(field, sample)
            # These should not be None
            self.assertIsNotNone(sample[field], f"Required field '{field}' is None")
            if field == "rationale":
                self.assertIsInstance(sample[field], str)
            elif field == "template_id":
                self.assertIsInstance(sample[field], int)
            elif field == "question_type":
                self.assertIsInstance(sample[field], str)
        
        self.logger.success("CoT fields tests passed")

    def test_example_data(self):
        """Print example data to show what the dataset looks like."""
        sample = self.train[0]
        self.logger.info("="*80)
        self.logger.info("EXAMPLE ECG-QA COT DATASET SAMPLE")
        self.logger.info("="*80)
        self.logger.info(f"Question: '{sample['question']}'")
        self.logger.info(f"Question type: '{sample['question_type']}'")
        self.logger.info(f"Template ID: {sample['template_id']}")
        self.logger.info(f"ECG IDs: {sample['ecg_id']}")
        self.logger.info(f"ECG paths: {sample['ecg_paths']}")
        if 'rationale' in sample:
            rationale_preview = sample['rationale'][:200] + "..." if len(sample['rationale']) > 200 else sample['rationale']
            self.logger.info(f"Rationale: '{rationale_preview}'")
        self.logger.info(f"Clinical contexts: {sample['clinical_contexts']}")
        self.logger.info("="*80)

class TestECGQACoTQADataset(unittest.TestCase):
    """
    Unit tests for the ECGQACoTQADataset class.
    """
    def setUp(self):
        # Set up global logger with verbose mode for detailed output
        set_global_verbose(True)
        self.logger = get_logger()
        
        from opentslm.time_series_datasets.ecg_qa import ECGQACoTQADataset
        self.ECGQACoTQADataset = ECGQACoTQADataset
        
        self.logger.loading("Initializing ECGQACoTQADataset...")
        # Use limited samples for faster testing
        self.train_dataset = self.ECGQACoTQADataset(split="train", EOS_TOKEN="", max_samples=5)
        self.val_dataset = self.ECGQACoTQADataset(split="validation", EOS_TOKEN="", max_samples=5)
        self.test_dataset = self.ECGQACoTQADataset(split="test", EOS_TOKEN="", max_samples=5)
        self.logger.success(f"Datasets initialized: Train={len(self.train_dataset)}, Val={len(self.val_dataset)}, Test={len(self.test_dataset)}")
        
        # Test the exclude_comparison functionality
        self.logger.loading("Testing exclude_comparison functionality...")
        self.train_no_comparison = self.ECGQACoTQADataset(split="train", EOS_TOKEN="", max_samples=10, exclude_comparison=True)
        self.logger.success(f"Non-comparison dataset initialized: {len(self.train_no_comparison)} samples")

    def test_dataset_sizes(self):
        """Test that the datasets are non-empty and splits are correct."""
        self.logger.info("Testing QA dataset sizes...")
        self.assertGreater(len(self.train_dataset), 0)
        self.assertGreater(len(self.val_dataset), 0)
        self.assertGreater(len(self.test_dataset), 0)
        self.logger.success("QA dataset size tests passed")

    def test_sample_keys(self):
        """Test that a sample contains all required keys."""
        self.logger.info("Testing QA sample keys...")
        sample = self.train_dataset[0]
        required_keys = {"answer", "pre_prompt", "post_prompt", "time_series", "time_series_text", "question", "question_type", "template_id"}
        self.assertTrue(required_keys.issubset(sample.keys()))
        
        # Test that required fields are not None
        for key in required_keys:
            self.assertIsNotNone(sample[key], f"Required field '{key}' is None")
        
        self.logger.success("QA sample keys test passed")

    def test_answer_is_rationale(self):
        """Test that the answer is a string (rationale)."""
        self.logger.info("Testing answer format...")
        sample = self.train_dataset[0]
        self.assertIsInstance(sample["answer"], str)
        self.assertGreater(len(sample["answer"]), 0)
        self.logger.success(f"Answer test passed: length={len(sample['answer'])}")

    def test_time_series_content(self):
        """Test that the time series and text are present and valid."""
        self.logger.info("Testing time series content...")
        sample = self.train_dataset[0]
        self.assertIsInstance(sample["time_series"], list)
        self.assertIsInstance(sample["time_series_text"], list)
        
        # Should have multiple time series (one per ECG lead)
        self.assertGreater(len(sample["time_series"]), 0)
        self.assertGreater(len(sample["time_series_text"]), 0)
        
        # Each time series should be non-empty
        for i, ts in enumerate(sample["time_series"]):
            self.assertGreater(len(ts), 0, f"Time series {i} is empty")
            # Check that all values are numeric
            for val in ts:
                self.assertIsInstance(val, (int, float))
        
        # Each time series text should be a string
        for i, ts_text in enumerate(sample["time_series_text"]):
            self.assertIsInstance(ts_text, str, f"Time series text {i} is not a string")
            self.assertGreater(len(ts_text), 0, f"Time series text {i} is empty")
        
        self.logger.success(f"Time series test passed: {len(sample['time_series'])} series")

    def test_time_series_text_includes_ecg_info(self):
        """Test that each time_series_text includes ECG lead information."""
        self.logger.info("Testing time series text format...")
        sample = self.train_dataset[0]
        for i, text in enumerate(sample['time_series_text']):
            self.logger.debug(f"Testing text {i}: {text[:100]}...")
            # Should mention ECG lead
            self.assertIn('ECG Lead', text)
            # Should mention mean and std
            self.assertIn('mean', text)
            self.assertIn('std', text)
            # Should be in natural language format (optimized version uses simpler format)
            # The optimized version uses: "This is ECG Lead X, it has mean Y and std Z:"
        
        self.logger.success("Time series text format tests passed")

    def test_labels_static_method(self):
        """Test that the static get_labels method returns expected labels."""
        self.logger.info("Testing static get_labels method...")
        labels = self.ECGQACoTQADataset.get_labels()
        
        # Test that the method works and returns reasonable results
        self.assertIsInstance(labels, list)
        self.assertGreater(len(labels), 0, "Labels list should not be empty")
        
        # Test that common labels are present
        common_labels = ["yes", "no", "none"]
        for label in common_labels:
            self.assertIn(label, labels, f"Common label '{label}' should be present")
        
        # Test that labels are strings
        for label in labels:
            self.assertIsInstance(label, str, f"All labels should be strings, got {type(label)}")
            self.assertGreater(len(label.strip()), 0, "Labels should not be empty strings")
        
        self.logger.success(f"Static labels test passed: {len(labels)} labels found")

    def test_prompts_content(self):
        """Test that prompts contain expected content."""
        self.logger.info("Testing prompt content...")
        sample = self.train_dataset[0]
        
        # Pre-prompt should mention cardiologist and ECG
        pre_prompt = sample["pre_prompt"]
        self.assertIn("cardiologist", pre_prompt.lower())
        self.assertIn("ecg", pre_prompt.lower())
        
        # Should mention the specific question
        question = sample["question"]
        self.assertIn(question, pre_prompt)
        
        # Post-prompt should mention answer format
        post_prompt = sample["post_prompt"]
        self.assertIn("Answer:", post_prompt)
        
        self.logger.success("Prompt content tests passed")

    def test_cot_specific_fields(self):
        """Test that CoT-specific fields are preserved in the formatted sample."""
        self.logger.info("Testing CoT-specific fields preservation...")
        sample = self.train_dataset[0]
        
        # Check fields that actually exist in the implementation
        cot_fields = ["rationale", "template_id", "question_type"]
        for field in cot_fields:
            if field in sample:
                self.logger.debug(f"Field {field}: {sample[field]}")
        
        self.logger.success("CoT fields preservation tests passed")

    def test_exclude_comparison_functionality(self):
        """Test that exclude_comparison properly filters out comparison questions."""
        self.logger.info("Testing exclude_comparison functionality...")
        
        # Check that the non-comparison dataset has no comparison questions
        for sample in self.train_no_comparison:
            question_type = sample.get("question_type")
            if question_type is None:
                raise ValueError(f"Sample missing question_type: {sample}")
            self.assertFalse(question_type.startswith("comparison"), 
                           f"Found comparison question in non-comparison dataset: {question_type}")
        
        # Check that the regular dataset might have comparison questions
        has_comparison = False
        for sample in self.train_dataset:
            question_type = sample.get("question_type")
            if question_type and question_type.startswith("comparison"):
                has_comparison = True
                break
        
        self.logger.success(f"Exclude comparison test passed. Regular dataset has comparison questions: {has_comparison}")

    def test_100hz_data_consistency(self):
        """Test that all time series data is consistently 100Hz."""
        self.logger.info("Testing 100Hz data consistency...")
        
        for dataset_name, dataset in [("train", self.train_dataset), ("val", self.val_dataset), ("test", self.test_dataset)]:
            for i, sample in enumerate(dataset):
                if i >= 3:  # Only test first 3 samples per dataset
                    break
                    
                time_series = sample.get("time_series")
                if time_series is None:
                    continue
                
                for j, ts in enumerate(time_series):
                    # Each time series should be exactly 1000 samples (10 seconds at 100Hz)
                    expected_length = 1000
                    actual_length = len(ts)
                    
                    if actual_length != expected_length:
                        self.logger.warning(f"Dataset {dataset_name}, sample {i}, time series {j}: "
                                          f"Expected {expected_length} samples, got {actual_length}")
                    
                    # All values should be numeric
                    for val in ts:
                        self.assertIsInstance(val, (int, float), 
                                           f"Non-numeric value in time series: {val}")
        
        self.logger.success("100Hz data consistency tests passed")

    def test_error_handling_missing_fields(self):
        """Test that the dataset properly raises errors for missing required fields."""
        self.logger.info("Testing error handling for missing fields...")
        
        # Create a sample with missing fields to test error handling
        from opentslm.time_series_datasets.ecg_qa import ECGQACoTQADataset
        
        # This should work normally
        try:
            normal_dataset = ECGQACoTQADataset(split="train", EOS_TOKEN="", max_samples=1)
            self.logger.success("Normal dataset creation works")
        except Exception as e:
            self.logger.error(f"Unexpected error creating normal dataset: {e}")
            raise
        
        # Test that we get proper errors when accessing samples with missing data
        # (This would happen if the underlying data is corrupted)
        self.logger.info("Error handling tests completed")

    def test_example_data_QA(self):
        """Print example data for ECGQACoTQADataset, showing all time series and text."""
        sample = self.train_dataset[0]
        self.logger.info("="*80)
        self.logger.info("EXAMPLE ECGQACoTQADataset SAMPLE")
        self.logger.info("="*80)
        self.logger.info(f"Question: '{sample['question']}'")
        self.logger.info(f"Question type: '{sample['question_type']}'")
        pre_prompt_preview = sample['pre_prompt'][:200]
        self.logger.info(f"Pre-prompt: '{pre_prompt_preview}'")
        self.logger.info(f"Post-prompt: '{sample['post_prompt']}'")
        answer_preview = sample['answer']
        self.logger.info(f"Answer (rationale): '{answer_preview}'")
        self.logger.info(f"Number of time series: {len(sample['time_series'])}")
        for i, (ts, ts_text) in enumerate(zip(sample['time_series'], sample['time_series_text'])):
            self.logger.info(f"Time series {i} text: '{ts_text}'")
            self.logger.info(f"Time series {i} length: {len(ts)}")
            self.logger.info(f"First 10 values: {ts[:10]}")
            self.logger.info(f"Last 10 values: {ts[-10:]}")
        self.logger.info("="*80)

if __name__ == "__main__":
    print("Running ECG-QA CoT loader and dataset tests, this might take a while...")
    unittest.main()