Fix pretty printing

test/util_test.py

@@ -1,7 +1,6 @@
 import unittest
 from typing import Iterator
-
-from . import test_data
+import xml.etree.ElementTree as ET
 
 from sia import util
 
@@ -115,4 +114,122 @@ class UtilTest(unittest.TestCase):
     def test_boolean_values(self):
         """Test boolean values are properly converted"""
         self.assertEqual(util.escape_text_for_xml(True), '<![CDATA[True]]>')
-        self.assertEqual(util.escape_text_for_xml(False), '<![CDATA[False]]>')
+        self.assertEqual(util.escape_text_for_xml(False), '<![CDATA[False]]>')
+
+    def test_pretty_print_element_basic(self):
+        """Test basic element printing"""
+        elem = ET.Element('root')
+        child = ET.SubElement(elem, 'child', {'attr': 'value'})
+        child.text = 'text'
+        
+        expected = '''<root>
+    <child attr="value">
+        <![CDATA[text]]>
+    </child>
+</root>'''
+        self.assertEqual(util.pretty_print_element(elem).strip(), expected.strip())
+
+    def test_pretty_print_element_long_attributes(self):
+        """Test printing elements with long attributes that need wrapping"""
+        elem = ET.Element('root', {
+            'very_long_attribute_1': 'value1',
+            'very_long_attribute_2': 'value2',
+            'very_long_attribute_3': 'value3'
+        })
+        
+        result = util.pretty_print_element(elem, max_line=40)
+        # Verify attributes are wrapped to new lines
+        self.assertIn('\n    very_long_attribute_1="value1"', result)
+        self.assertIn('\n    very_long_attribute_2="value2"', result)
+        self.assertIn('\n    very_long_attribute_3="value3"', result)
+
+    def test_pretty_print_element_cdata(self):
+        """Test printing elements with CDATA sections"""
+        elem = ET.Element('root')
+        elem.text = 'preserved\n spacing'
+        
+        expected = '''<root>
+    <![CDATA[preserved\n spacing]]>
+</root>'''
+        self.assertEqual(util.pretty_print_element(elem).strip(), expected.strip())
+
+    def test_pretty_print_element_empty(self):
+        """Test printing empty elements"""
+        elem = ET.Element('empty')
+        self.assertEqual(util.pretty_print_element(elem), '<empty/>')
+
+    def test_real_world_context(self):
+        """Test formatting a realistic context example"""
+        # Create a context element with system metrics and a repeat entry
+        context = ET.Element('context', {
+            'time': '2024-10-18T12:00:00Z',
+            'cpu': '12',
+            'gpu': '26',
+            'memory_used': '9556302234',
+            'memory_total': '17179869184',
+            'disk_used': '244434939904',
+            'disk_total': '273145991168',
+            'context': '3',
+            'stdin': '0'
+        })
+    
+        repeat = ET.SubElement(context, 'repeat', {
+            'id': 'a3d89ee5-28ec-4c5a-b9e9-a30af53d43a0',
+            'exit_code': '0'
+        })
+        repeat.text = 'ls -lah /'
+    
+        stdout = ET.SubElement(repeat, 'stdout')
+        stdout.text = '''total 16K
+drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 ./
+drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 ../
+drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 tasks/
+drwxr-xr-x 1 sia 1049089 0 Oct 28 13:40 user/'''
+    
+        stderr = ET.SubElement(repeat, 'stderr')
+    
+        formatted = util.pretty_print_element(context)
+    
+        # Verify key formatting expectations
+        self.assertIn('    <repeat exit_code="0" id="a3d89ee5-28ec-4c5a-b9e9-a30af53d43a0">', formatted)
+        self.assertIn('        <![CDATA[ls -lah /]]>', formatted)
+        self.assertIn('        <stdout>', formatted)
+        self.assertIn('            <![CDATA[total 16K', formatted)
+
+    def test_large_content(self):
+        """Test formatting with large content blocks"""
+        # Create element with large text content
+        elem = ET.Element('root')
+        elem.text = 'x' * 1000
+        
+        result = util.pretty_print_element(elem)
+        
+        # Verify content is properly indented and wrapped
+        self.assertIn('x' * 1000, result)
+
+    def test_special_characters(self):
+        """Test handling of special XML characters"""
+        elem = ET.Element('root')
+        child = ET.SubElement(elem, 'child')
+        child.text = '< > & " \''
+        
+        result = util.pretty_print_element(elem)
+        
+        # Should be escaped or wrapped in CDATA
+        self.assertTrue(
+            ('&lt; &gt; &amp; &quot; &apos;' in result) or
+            ('<![CDATA[< > & " \']]>' in result)
+        )
+
+    def test_mixed_content(self):
+        """Test formatting elements with mixed text and child nodes"""
+        elem = ET.Element('root')
+        elem.text = 'before'
+        child = ET.SubElement(elem, 'child')
+        child.tail = 'after'
+        
+        result = util.pretty_print_element(elem)
+        
+        # Verify text positioning
+        self.assertIn('before', result)
+        self.assertIn('after', result)