import unittest
import dis
from opcode import *
class TestFindLabels(unittest.TestCase):
	# Since the behaviour of the jump opcodes is mostly similiar
	# With the exception of JUMP_FORWARD, there is no point writing 
	# separate tests for all of them. 
	def test_original(self):
		sample_code = bytes(
		    chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(0) +
		    chr(EXTENDED_ARG) + chr(1) + chr(0) +
		    chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(0) +
		    chr(opmap["RETURN_VALUE"]),
		    encoding="latin-1"
		)
		expected_output = [3, 65545]
		self.assertEqual(dis.findlabels(sample_code), expected_output)
	def test_jump_if_false_or_pop(self):
		sample_code = bytes (
			chr(opmap["JUMP_IF_FALSE_OR_POP"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(1) + chr(0) + 
			chr(opmap["JUMP_IF_FALSE_OR_POP"]) + chr(0) + chr(1) +
			chr(EXTENDED_ARG) + chr(2) + chr(0) +
			chr(opmap["JUMP_IF_FALSE_OR_POP"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(255) + chr(255) +
			chr(opmap["JUMP_IF_FALSE_OR_POP"]) + chr(0) + chr(0) +
			chr(opmap["RETURN_VALUE"]),
			encoding="latin-1"
			)
		expected_output = [0, 65792, 131072, 4294901760]
		self.assertEqual(dis.findlabels(sample_code), expected_output)
	def test_jump_forward(self):
		sample_code = bytes (
			chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(1) + chr(0) + 
			chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(1) +
			chr(EXTENDED_ARG) + chr(2) + chr(0) +
			chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(255) + chr(255) +
			chr(opmap["JUMP_FORWARD"]) + chr(0) + chr(0) +
			chr(opmap["RETURN_VALUE"]),
			encoding="latin-1"
			)
		expected_output = [3, 65801, 131087, 4294901781]
		self.assertEqual(dis.findlabels(sample_code), expected_output)
	def test_jump_absolute(self):
		sample_code = bytes (
			chr(opmap["JUMP_ABSOLUTE"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(1) + chr(0) + 
			chr(opmap["JUMP_ABSOLUTE"]) + chr(0) + chr(1) +
			chr(EXTENDED_ARG) + chr(2) + chr(0) +
			chr(opmap["JUMP_ABSOLUTE"]) + chr(0) + chr(0) +
			chr(EXTENDED_ARG) + chr(255) + chr(255) +
			chr(opmap["JUMP_ABSOLUTE"]) + chr(0) + chr(0) +
			chr(opmap["RETURN_VALUE"]),
			encoding="latin-1"
			)
		expected_output = [0, 65792, 131072, 4294901760]
		self.assertEqual(dis.findlabels(sample_code), expected_output)

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