import unittest
import dis
from opcode import *


class TestFindLabels(unittest.TestCase):
    def test_findlabels(self):
        args = (
            0x0000,     # Unextended argument values.
            0xFFFF,
            0x00010100, # Extended arguments.  Note: byte swapped from ``sample_code``!
            0x00020000,
            0xFFFF0000,
            0xFFFFFFFF,
            )
        offsets = (     # "Program counter" after each jump instruction. 
            3,
            6,
            12,
            18,
            24,
            30,
            )

        non_jump_labels = []
        abs_jump_labels = list(args)
        rel_jump_labels = [arg + ofs for arg, ofs in zip(args, offsets)]

        known_non_arg_opcode = opmap["RETURN_VALUE"]
        for opcode in range(256):
            if opcode >= HAVE_ARGUMENT:
                sample_code = bytes([
                    opcode,       0x00, 0x00,
                    opcode,       0xFF, 0xFF,
                    EXTENDED_ARG, 0x01, 0x00, opcode, 0x00, 0x01,
                    EXTENDED_ARG, 0x02, 0x00, opcode, 0x00, 0x00,
                    EXTENDED_ARG, 0xFF, 0xFF, opcode, 0x00, 0x00,
                    EXTENDED_ARG, 0xFF, 0xFF, opcode, 0xFF, 0xFF,
                    known_non_arg_opcode,
                    ])
                labels = dis.findlabels(sample_code)
                if opcode in hasjrel:
                    self.assertEqual(labels, rel_jump_labels)
                elif opcode in hasjabs:
                    self.assertEqual(labels, abs_jump_labels)
                else:
                    self.assertEqual(labels, non_jump_labels)

                    
suite = unittest.TestLoader().loadTestsFromTestCase(TestFindLabels)
unittest.TextTestRunner(verbosity=2).run(suite)