Diff of /tests/test_preprocess.py [000000] .. [8c54ae]

Switch to side-by-side view

--- a
+++ b/tests/test_preprocess.py
@@ -0,0 +1,83 @@
+import sys
+import os
+import unittest
+import xml.etree.ElementTree as ET
+from src.ai.preprocess import extract_section_data
+
+# Add the parent directory to the system path for module imports
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+class TestExtractSectionData(unittest.TestCase):
+    """
+    Unit tests for the extract_section_data function in the preprocessing module.
+
+    This class contains test cases to validate the extraction of medication 
+    information from XML data.
+    """
+
+    def setUp(self):
+        """
+        Set up test environment by creating sample XML data and initializing 
+        patient data dictionary before each test.
+        """
+        # Sample XML data for testing
+        self.xml_data = '''<root xmlns:hl7="urn:hl7-org:v3">
+            <hl7:component>
+                <hl7:section>
+                    <hl7:title>Medications</hl7:title>
+                    <hl7:tbody>
+                        <hl7:tr>
+                            <hl7:td>2023-01-01T00:00:00Z</hl7:td>
+                            <hl7:td>2023-01-10T00:00:00Z</hl7:td>
+                            <hl7:td>Aspirin</hl7:td>
+                        </hl7:tr>
+                        <hl7:tr>
+                            <hl7:td>2023-02-01T00:00:00Z</hl7:td>
+                            <hl7:td>2023-02-05T00:00:00Z</hl7:td>
+                            <hl7:td>Ibuprofen</hl7:td>
+                        </hl7:tr>
+                    </hl7:tbody>
+                </hl7:section>
+            </hl7:component>
+        </root>'''
+        # Parse the XML string into an ElementTree
+        self.root = ET.fromstring(self.xml_data)
+        # Initialize a dictionary to hold patient data
+        self.patient_data = {}
+
+    def test_extract_medications(self):
+        """
+        Test the extraction of medication data from the XML structure.
+
+        This test checks whether the extract_section_data function correctly 
+        populates the patient_data dictionary with medication information.
+        """
+        # Call the function to extract medication data
+        extract_section_data("Medications", self.root.find('.//hl7:section', {'hl7': 'urn:hl7-org:v3'}), self.patient_data)
+
+        # Define the expected data after extraction
+        expected_data = [
+            {
+                'Start': '2023-01-01T00:00:00Z',
+                'Stop': '2023-01-10T00:00:00Z',
+                'Description': 'Aspirin',
+                'Duration of Usage': '10 days',  # Assumed that calculate_duration works correctly
+                'Last Usage': '636 days ago'  # Replace with expected value based on your logic
+            },
+            {
+                'Start': '2023-02-01T00:00:00Z',
+                'Stop': '2023-02-05T00:00:00Z',
+                'Description': 'Ibuprofen',
+                'Duration of Usage': '5 days',  # Assumed that calculate_duration works correctly
+                'Last Usage': '610 days ago'  # Replace with expected value based on your logic
+            }
+        ]
+
+        # Validate the data extracted
+        self.assertIn('Medications', self.patient_data)  # Check if 'Medications' key exists
+        self.assertEqual(len(self.patient_data['Medications']), 2)  # Validate number of entries
+        self.assertEqual(self.patient_data['Medications'], expected_data)  # Check if extracted data matches expected
+
+# Entry point for running the unit tests
+if __name__ == '__main__':
+    unittest.main()