

from unittest import TestCase
from tlexport.checksums import ones_complement_checksum
from tlexport.checksums import calculate_checksum_tcp
from tlexport.packet import Packet


class TestChecksums(TestCase):
    def setUp(self) -> None:
        self.string_1 = bytearray(
            b'\x7f\x00\x00\x01\x7f\x00\x00\x01\x00\x06\x00('
            b'\x8f\xfe\xad*\xaf;\xe7\xac\x00\x00\x00\x00\xa0\x02\xff\xd7\x00\x00\x00\x00\x02\x04\xff\xd7\x04\x02\x08'
            b'\nX\xf7`\xee\x00\x00\x00\x00\x01\x03\x03\x07')
        self.string_2 = bytearray(
            b'\x7f\x00\x00\x01\x7f\x00\x00\x01\x00\x06\x00('
            b'\xad*\x8f\xfe\xd7\xcfN\xb8\xaf;\xe7\xad\xa0\x12\xff\xcb\x00\x00\x00\x00\x02\x04\xff\xd7\x04\x02\x08\nX'
            b'\xf7`\xeeX\xf7`\xee\x01\x03\x03\x07')

        self.false_string_1 = bytearray(
            b'\x7f\x00\x00\x01\x7f\x00\x00\x01\x00\x06\x00('
            b'\x8f\xfe\xad*\xaf;\xe7\xac\x00\x00\x00\x00\xa0\x02\xff\xd7\x00\x00\x00\x00\x02\x04\xff\xd7')

        self.zeroes = bytearray(b'')

        self.packet_1 = Packet(
            b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08\x00E\x00\x00<\x11\xdc@\x00@\x06*\xde\x7f\x00\x00'
            b'\x01\x7f\x00\x00\x01\x8f\xfe\xad*\xaf;\xe7\xac\x00\x00\x00\x00\xa0\x02\xff\xd7\xc2\x0a\x00\x00\x02\x04'
            b'\xff\xd7\x04\x02\x08\nX\xf7`\xee\x00\x00\x00\x00\x01\x03\x03\x07',
            1693392019.4963138)
        self.bad_packet_1 = Packet(
            b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08\x00E\x00\x00<\x00\x00@\x00@\x06<\xba\x7f\x00\x00'
            b'\x01\x7f\x00\x00\x01\xad*\x8f\xfe\xd7\xcfN\xb8\xaf;\xe7\xad\xa0\x12\xff\xcb\xfe0\x00\x00\x02\x04\xff'
            b'\xd7\x04\x02\x08\nX\xf7`\xeeX\xf7`\xee\x01\x03\x03\x07',
            1693392019.4963367)

    def test_ones_complement_checksum(self):
        self.assertEqual(ones_complement_checksum(self.string_1), bytearray.fromhex("c20a"))
        self.assertEqual(ones_complement_checksum(self.string_2), bytearray.fromhex("e197"))
        self.assertNotEqual(ones_complement_checksum(self.false_string_1), bytearray.fromhex("ab12"))

        self.assertEqual(ones_complement_checksum(self.zeroes), bytearray.fromhex("ffff"))

    def test_tcp_checksum(self):
        self.assertEqual(calculate_checksum_tcp(self.packet_1), True)
        self.assertEqual(calculate_checksum_tcp(self.bad_packet_1), False)
