1 files changed, 50 insertions, 26 deletions
diff --git a/tests/arch/operands/immediate.py b/tests/arch/operands/immediate.py
index 74b8069..c3fcb84 100644
--- a/tests/arch/operands/immediate.py
+++ b/tests/arch/operands/immediate.py
@@ -1,35 +1,59 @@
-#!/usr/bin/python3-dbg
-# -*- coding: utf-8 -*-
-
 
 import pychrysalide
-from chrysacase import ChrysalideTestCase
-from pychrysalide import arch
-from pychrysalide.arch import ImmOperand
 
+from chrysacase import ChrysalideTestCase
+from pychrysalide import MemoryDataSize
+from pychrysalide.arch.operands import ImmediateOperand
+from pychrysalide.glibext import StringBuilder
 
 
 class TestImmediate(ChrysalideTestCase):
-    """TestCase for arch.ImmOperand."""
+    """TestCase for arch.ImmediateOperand."""
+
+
+    def testBasicImmediate(self):
+        """Check basic properties of immediate values."""
+
+        imm = ImmediateOperand(MemoryDataSize._32_BITS_UNSIGNED, 0x123)
+
+        self.assertEqual(imm.size, MemoryDataSize._32_BITS_UNSIGNED)
+        self.assertEqual(imm.value, 0x123)
+        self.assertFalse(imm.is_negative)
+
+
+    def testStringBuilderForImmediatesOverriding(self):
+        """Override the StringBuilder interface for immediate values."""
+
+        class MyImmOperand(ImmediateOperand, StringBuilder):
+
+            def __init__(self):
+                super().__init__(MemoryDataSize._32_BITS_UNSIGNED, 0x123)
+
+            def _to_string(self, flags=0):
+                return 'NaN'
+
+        op = MyImmOperand()
+
+        self.assertEqual(op.to_string(), 'NaN')
+        self.assertEqual(str(op), 'NaN')
+        self.assertEqual(f'{op}', 'NaN')
 
 
     def validateValue(self, value, size, padding, strings):
         """Check all kinds of things with a given immediate operand."""
 
-        display = [
-            ImmOperand.IOD_BIN, ImmOperand.IOD_OCT,
-            ImmOperand.IOD_DEC,
-            ImmOperand.IOD_HEX
-        ]
+        for d in strings.keys():
 
-        for d in display:
-
-            op = ImmOperand(size, value)
+            op = ImmediateOperand(size, value)
 
             self.assertTrue(op.size == size)
             self.assertTrue(op.value == value)
 
-            op.padding = padding
+            if padding:
+                op.set_flag(ImmediateOperand.ImmOperandFlag.ZERO_PADDING)
+            else:
+                op.unset_flag(ImmediateOperand.ImmOperandFlag.ZERO_PADDING_BY_DEFAULT)
+
             op.display = d
 
             string = op.to_string()
@@ -40,23 +64,23 @@ class TestImmediate(ChrysalideTestCase):
         """Run sanity checks on immediate operand with value 1."""
 
         strings = {
-            ImmOperand.IOD_BIN: 'b1',
-            ImmOperand.IOD_OCT: '01',
-            ImmOperand.IOD_DEC: '1',
-            ImmOperand.IOD_HEX: '0x1'
+            ImmediateOperand.ImmOperandDisplay.BIN: 'b1',
+            ImmediateOperand.ImmOperandDisplay.OCT: '01',
+            ImmediateOperand.ImmOperandDisplay.DEC: '1',
+            ImmediateOperand.ImmOperandDisplay.HEX: '0x1'
         }
 
-        self.validateValue(1, arch.MDS_8_BITS_UNSIGNED, False, strings)
+        self.validateValue(1, pychrysalide.MemoryDataSize._8_BITS_UNSIGNED, False, strings)
 
 
     def testByteOnePadded(self):
         """Run sanity checks on immediate operand with padded value 1."""
 
         strings = {
-            ImmOperand.IOD_BIN: 'b00000001',
-            ImmOperand.IOD_OCT: '01',
-            ImmOperand.IOD_DEC: '1',
-            ImmOperand.IOD_HEX: '0x01'
+            ImmediateOperand.ImmOperandDisplay.BIN: 'b00000001',
+            ImmediateOperand.ImmOperandDisplay.OCT: '01',
+            ImmediateOperand.ImmOperandDisplay.DEC: '1',
+            ImmediateOperand.ImmOperandDisplay.HEX: '0x01'
         }
 
-        self.validateValue(1, arch.MDS_8_BITS_UNSIGNED, True, strings)
+        self.validateValue(1, pychrysalide.MemoryDataSize._8_BITS_UNSIGNED, True, strings)