src/target/trx_toolkit/test_codec.py

   1 #!/usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 '''
   5 Unit tests for declarative message codec.
   6 '''
   7
   8 # (C) 2021 by sysmocom - s.f.m.c. GmbH <info@sysmocom.de>
   9 # Author: Vadim Yanitskiy <vyanitskiy@sysmocom.de>
  10 #
  11 # All Rights Reserved
  12 #
  13 # This program is free software; you can redistribute it and/or modify
  14 # it under the terms of the GNU General Public License as published by
  15 # the Free Software Foundation; either version 2 of the License, or
  16 # (at your option) any later version.
  17 #
  18 # This program is distributed in the hope that it will be useful,
  19 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  20 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  21 # GNU General Public License for more details.
  22
  23 import unittest
  24 import struct
  25 import codec
  26
  27 class TestField(codec.Field):
  28         DEF_PARAMS = { 'key' : 0xde }
  29         DEF_LEN = 4
  30
  31         @staticmethod
  32         def xor(data: bytes, key: int = 0x00):
  33                 return bytes([x ^ key for x in data])
  34
  35         def _from_bytes(self, vals: dict, data: bytes) -> None:
  36                 vals[self.name] = self.xor(data, self.p['key'])
  37
  38         def _to_bytes(self, vals: dict) -> bytes:
  39                 return self.xor(self.get_val(vals), self.p['key'])
  40
  41 class Field(unittest.TestCase):
  42         MAGIC = b'\xde\xad\xbe\xef'
  43
  44         def test_to_bytes(self):
  45                 vals = { 'magic' : self.MAGIC, 'other' : 'unrelated' }
  46                 encoded_de = TestField.xor(self.MAGIC, 0xde)
  47                 encoded_88 = TestField.xor(self.MAGIC, 0x88)
  48
  49                 with self.subTest('default length=4, default key=0xde'):
  50                         field = TestField('magic')
  51                         self.assertEqual(field.to_bytes(vals), encoded_de)
  52
  53                 with self.subTest('default length=4, different key=0x88'):
  54                         field = TestField('magic', key=0x88)
  55                         self.assertEqual(field.to_bytes(vals), encoded_88)
  56
  57                 with self.subTest('different length=2, default key=0xde'):
  58                         field = TestField('magic', len=2)
  59                         vals['magic'] = vals['magic'][:2]
  60                         self.assertEqual(field.to_bytes(vals), encoded_de[:2])
  61
  62                 with self.subTest('EncodeError due to length mismatch'):
  63                         field = TestField('magic', len=8)
  64                         with self.assertRaises(codec.EncodeError):
  65                                 field.to_bytes(vals)
  66
  67         def test_from_bytes(self):
  68                 encoded_de = TestField.xor(self.MAGIC, 0xde) + b'\xff' * 60
  69                 encoded_88 = TestField.xor(self.MAGIC, 0x88) + b'\xff' * 60
  70                 vals = { 'magic' : 'overrien', 'other' : 'unchanged' }
  71
  72                 with self.subTest('default length=4, default key=0xde'):
  73                         field = TestField('magic')
  74                         offset = field.from_bytes(vals, encoded_de)
  75                         self.assertEqual(vals['other'], 'unchanged')
  76                         self.assertEqual(vals['magic'], self.MAGIC)
  77                         self.assertEqual(offset, len(self.MAGIC))
  78
  79                 with self.subTest('default length=4, different key=0x88'):
  80                         field = TestField('magic', key=0x88)
  81                         offset = field.from_bytes(vals, encoded_88)
  82                         self.assertEqual(vals['other'], 'unchanged')
  83                         self.assertEqual(vals['magic'], self.MAGIC)
  84                         self.assertEqual(offset, len(self.MAGIC))
  85
  86                 with self.subTest('different length=2, default key=0xde'):
  87                         field = TestField('magic', len=2)
  88                         offset = field.from_bytes(vals, encoded_de)
  89                         self.assertEqual(vals['other'], 'unchanged')
  90                         self.assertEqual(vals['magic'], self.MAGIC[:2])
  91                         self.assertEqual(offset, 2)
  92
  93                 with self.subTest('full length, different key=0x88'):
  94                         field = TestField('magic', len=0, key=0x88)
  95                         offset = field.from_bytes(vals, encoded_88)
  96                         self.assertEqual(vals['other'], 'unchanged')
  97                         self.assertEqual(vals['magic'], self.MAGIC + b'\x77' * 60)
  98                         self.assertEqual(offset, len(encoded_88))
  99
 100                 with self.subTest('DecodeError due to short read'):
 101                         field = TestField('magic', len=4)
 102                         with self.assertRaises(codec.DecodeError):
 103                                 field.from_bytes(vals, b'\x00')
 104
 105         def test_get_pres(self):
 106                 vals = { 'magic' : self.MAGIC }
 107
 108                 with self.subTest('to_bytes() for a non-existing field'):
 109                         field = TestField('not-there')
 110                         with self.assertRaises(KeyError):
 111                                 field.to_bytes(vals)
 112
 113                 with self.subTest('to_bytes() for a field with get_pres()'):
 114                         field = TestField('magic', key=0x00)
 115                         field.get_pres = lambda v: not v['omit']
 116
 117                         data = field.to_bytes({ **vals, 'omit' : False })
 118                         self.assertEqual(data, self.MAGIC)
 119
 120                         data = field.to_bytes({ **vals, 'omit' : True })
 121                         self.assertEqual(data, b'')
 122
 123                 with self.subTest('from_bytes() for a field with get_pres()'):
 124                         field = TestField('magic', key=0x00)
 125                         field.get_pres = lambda v: not v['omit']
 126
 127                         vals = { 'omit' : False }
 128                         offset = field.from_bytes(vals, self.MAGIC)
 129                         self.assertEqual(vals['magic'], self.MAGIC)
 130                         self.assertEqual(offset, len(self.MAGIC))
 131
 132                         vals = { 'omit' : True }
 133                         offset = field.from_bytes(vals, self.MAGIC)
 134                         self.assertFalse('magic' in vals)
 135                         self.assertEqual(offset, 0)
 136
 137         def test_get_len(self):
 138                 vals = { 'len' : 32, 'unrelated' : 'foo' }
 139
 140                 field = TestField('magic', key=0x00)
 141                 field.get_len = lambda v, _: v['len']
 142
 143                 with self.subTest('not enough octets in the buffer: 16 < 32'):
 144                         with self.assertRaises(codec.DecodeError):
 145                                 field.from_bytes(vals, b'\xff' * 16)
 146
 147                 with self.subTest('more than enough octets in the buffer'):
 148                         offset = field.from_bytes(vals, b'\xff' * 64)
 149                         self.assertEqual(vals['magic'], b'\xff' * 32)
 150                         self.assertEqual(offset, 32)
 151
 152                 with self.subTest('length field does not exist'):
 153                         with self.assertRaises(KeyError):
 154                                 field.from_bytes({ }, b'\xff' * 64)
 155
 156         def test_get_val(self):
 157                 field = TestField('magic', key=0x00, len=0)
 158                 field.get_val = lambda v: v.get('val', self.MAGIC)
 159
 160                 with self.subTest('value is present in the dict'):
 161                         data = field.to_bytes({ 'val' : b'\xd0\xde' })
 162                         self.assertEqual(data, b'\xd0\xde')
 163
 164                 with self.subTest('value is not present in the dict'):
 165                         data = field.to_bytes({ })
 166                         self.assertEqual(data, self.MAGIC)
 167
 168 class Buf(unittest.TestCase):
 169         MAGIC = b'\xde\xad' * 4
 170
 171         def test_to_bytes(self):
 172                 vals = { 'buf' : self.MAGIC }
 173
 174                 with self.subTest('with no length constraints'):
 175                         field = codec.Buf('buf') # default: len=0
 176                         self.assertEqual(field.to_bytes(vals), self.MAGIC)
 177
 178                 with self.subTest('with length constraints'):
 179                         field = codec.Buf('buf', len=len(self.MAGIC))
 180                         self.assertEqual(field.to_bytes(vals), self.MAGIC)
 181
 182                 with self.subTest('EncodeError due to length mismatch'):
 183                         field = codec.Buf('buf', len=4)
 184                         with self.assertRaises(codec.EncodeError):
 185                                 field.to_bytes(vals)
 186
 187         def test_from_bytes(self):
 188                 vals = { }
 189
 190                 with self.subTest('with no length constraints'):
 191                         field = codec.Buf('buf') # default: len=0
 192                         offset = field.from_bytes(vals, self.MAGIC)
 193                         self.assertEqual(vals['buf'], self.MAGIC)
 194                         self.assertEqual(offset, len(self.MAGIC))
 195
 196                 with self.subTest('with length constraints'):
 197                         field = codec.Buf('buf', len=2)
 198                         offset = field.from_bytes(vals, self.MAGIC)
 199                         self.assertEqual(vals['buf'], self.MAGIC[:2])
 200                         self.assertEqual(offset, len(self.MAGIC[:2]))
 201
 202                 with self.subTest('DecodeError due to not enough bytes'):
 203                         field = codec.Buf('buf', len=64)
 204                         with self.assertRaises(codec.DecodeError):
 205                                 field.from_bytes(vals, self.MAGIC)
 206
 207 class Spare(unittest.TestCase):
 208         # Fixed length with custom filler
 209         SAA = codec.Spare('pad', len=4, filler=b'\xaa')
 210         # Auto-calculated length with custom filler
 211         SFF = codec.Spare('pad', filler=b'\xff')
 212         SFF.get_len = lambda v, _: v['len']
 213         # Fixed length with default filler
 214         S00 = codec.Spare('pad', len=2)
 215
 216         def test_to_bytes(self):
 217                 self.assertEqual(self.SFF.to_bytes({ 'len' : 8 }), b'\xff' * 8)
 218                 self.assertEqual(self.SAA.to_bytes({ }), b'\xaa' * 4)
 219                 self.assertEqual(self.S00.to_bytes({ }), b'\x00' * 2)
 220
 221         def test_from_bytes(self):
 222                 with self.assertRaises(codec.DecodeError):
 223                         self.S00.from_bytes({ }, b'\x00') # Short read
 224                 self.assertEqual(self.SFF.from_bytes({ 'len' : 8 }, b'\xff' * 8), 8)
 225                 self.assertEqual(self.SAA.from_bytes({ }, b'\xaa' * 64), 4)
 226                 self.assertEqual(self.S00.from_bytes({ }, b'\x00' * 64), 2)
 227
 228 class Uint(unittest.TestCase):
 229         def _test_uint(self, field, fmt, vals):
 230                 for i in vals:
 231                         with self.subTest('to_bytes()'):
 232                                 val = field.to_bytes({ field.name : i })
 233                                 self.assertEqual(val, struct.pack(fmt, i))
 234
 235                         with self.subTest('from_bytes()'):
 236                                 data, parsed = struct.pack(fmt, i), { }
 237                                 offset = field.from_bytes(parsed, data)
 238                                 self.assertEqual(offset, len(data))
 239                                 self.assertEqual(parsed[field.name], i)
 240
 241         def test_uint8(self):
 242                 self._test_uint(codec.Uint('foo'), 'B', range(2 ** 8))
 243
 244         def test_int8(self):
 245                 self._test_uint(codec.Int('foo'), 'b', range(-128, 128))
 246
 247         def test_uint16(self):
 248                 vals = (0, 65, 128, 255, 512, 1023, 2 ** 16 - 1)
 249                 self._test_uint(codec.Uint16BE('foo'), '>H', vals)
 250                 self._test_uint(codec.Uint16LE('foo'), '<H', vals)
 251
 252         def test_int16(self):
 253                 vals = (-32767, -16384, 0, 16384, 32767)
 254                 self._test_uint(codec.Int16BE('foo'), '>h', vals)
 255                 self._test_uint(codec.Int16LE('foo'), '<h', vals)
 256
 257         def test_uint32(self):
 258                 vals = (0, 33, 255, 1024, 1337, 4099, 2 ** 32 - 1)
 259                 self._test_uint(codec.Uint32BE('foo'), '>I', vals)
 260                 self._test_uint(codec.Uint32LE('foo'), '<I', vals)
 261
 262         def test_int32(self):
 263                 vals = (-2147483647, 0, 2147483647)
 264                 self._test_uint(codec.Int32BE('foo'), '>i', vals)
 265                 self._test_uint(codec.Int32LE('foo'), '<i', vals)
 266
 267         def test_offset_mult(self):
 268                 with self.subTest('encode / decode with offset=5'):
 269                         field = codec.Uint('foo', offset=5)
 270
 271                         self.assertEqual(field.to_bytes({ 'foo' : 10 }), b'\x05')
 272                         self.assertEqual(field.to_bytes({ 'foo' :  5 }), b'\x00')
 273
 274                         vals = { 'foo' : 'overriden' }
 275                         field.from_bytes(vals, b'\xff')
 276                         self.assertEqual(vals['foo'], 260)
 277                         field.from_bytes(vals, b'\x00')
 278                         self.assertEqual(vals['foo'], 5)
 279
 280                 with self.subTest('encode / decode with mult=2'):
 281                         field = codec.Uint('foo', mult=2)
 282
 283                         self.assertEqual(field.to_bytes({ 'foo' : 0 }), b'\x00')
 284                         self.assertEqual(field.to_bytes({ 'foo' : 3 }), b'\x01')
 285                         self.assertEqual(field.to_bytes({ 'foo' : 32 }), b'\x10')
 286                         self.assertEqual(field.to_bytes({ 'foo' : 64 }), b'\x20')
 287
 288                         vals = { 'foo' : 'overriden' }
 289                         field.from_bytes(vals, b'\x00')
 290                         self.assertEqual(vals['foo'], 0 * 2)
 291                         field.from_bytes(vals, b'\x0f')
 292                         self.assertEqual(vals['foo'], 15 * 2)
 293                         field.from_bytes(vals, b'\xff')
 294                         self.assertEqual(vals['foo'], 255 * 2)
 295
 296 class BitFieldSet(unittest.TestCase):
 297         S16 = codec.BitFieldSet(set=(
 298                 codec.BitField('f4a', bl=4),
 299                 codec.BitField('f8', bl=8),
 300                 codec.BitField('f4b', bl=4),
 301         ))
 302
 303         S8M = codec.BitFieldSet(order='msb', set=(
 304                 codec.BitField('f4', bl=4),
 305                 codec.BitField('f1', bl=1),
 306                 codec.BitField('f3', bl=3),
 307         ))
 308
 309         S8L = codec.BitFieldSet(order='lsb', set=(
 310                 codec.BitField('f4', bl=4),
 311                 codec.BitField('f1', bl=1),
 312                 codec.BitField('f3', bl=3),
 313         ))
 314
 315         S8V = codec.BitFieldSet(set=(
 316                 codec.BitField('f4', bl=4, val=2),
 317                 codec.BitField('f1', bl=1, val=0),
 318                 codec.BitField('f3', bl=3),
 319         ))
 320
 321         S8P = codec.BitFieldSet(set=(
 322                 codec.BitField.Spare(bl=4),
 323                 codec.BitField('f4', bl=4),
 324         ))
 325
 326         @staticmethod
 327         def from_bytes(s: codec.BitFieldSet, data: bytes) -> dict:
 328                 vals = { }
 329                 s.from_bytes(vals, data)
 330                 return vals
 331
 332         def test_len_auto(self):
 333                 with self.subTest('1 + 2 = 3 bits => 1 octet (with padding)'):
 334                         s = codec.BitFieldSet(set=(
 335                                 codec.BitField('f1', bl=1),
 336                                 codec.BitField('f2', bl=2),
 337                         ))
 338                         self.assertEqual(s.len, 1)
 339
 340                 with self.subTest('4 + 2 + 2 = 8 bits => 1 octet'):
 341                         s = codec.BitFieldSet(set=(
 342                                 codec.BitField('f4', bl=4),
 343                                 codec.BitField('f2a', bl=2),
 344                                 codec.BitField('f2b', bl=2),
 345                         ))
 346                         self.assertEqual(s.len, 1)
 347
 348                 with self.subTest('12 + 4 + 2 = 18 bits => 3 octets (with padding)'):
 349                         s = codec.BitFieldSet(set=(
 350                                 codec.BitField('f12', bl=12),
 351                                 codec.BitField('f4', bl=4),
 352                                 codec.BitField('f2', bl=2),
 353                         ))
 354                         self.assertEqual(s.len, 3)
 355
 356         def test_overflow(self):
 357                 with self.assertRaises(codec.ProtocolError):
 358                         s = codec.BitFieldSet(len=1, set=(
 359                                 codec.BitField('f6', bl=6),
 360                                 codec.BitField('f4', bl=4),
 361                         ))
 362
 363         def test_offset_mask(self):
 364                 calc = lambda s: [(f.name, f.offset, f.mask) for f in s._fields]
 365
 366                 with self.subTest('16 bit total (MSB): f4a + f8 + f4b'):
 367                         om = [('f4a', 8 + 4, 0x0f), ('f8', 4, 0xff), ('f4b', 0, 0x0f)]
 368                         self.assertEqual(len(self.S16._fields), 3)
 369                         self.assertEqual(calc(self.S16), om)
 370
 371                 with self.subTest('8 bit total (MSB): f4 + f1 + f3'):
 372                         om = [('f4', 1 + 3, 0x0f), ('f1', 3, 0x01), ('f3', 0, 0x07)]
 373                         self.assertEqual(len(self.S8M._fields), 3)
 374                         self.assertEqual(calc(self.S8M), om)
 375
 376                 with self.subTest('8 bit total (LSB): f4 + f1 + f3'):
 377                         om = [('f3', 1 + 4, 0x07), ('f1', 4, 0x01), ('f4', 0, 0x0f)]
 378                         self.assertEqual(len(self.S8L._fields), 3)
 379                         self.assertEqual(calc(self.S8L), om)
 380
 381                 with self.subTest('8 bit total (LSB): s4 + f4'):
 382                         om = [(None, 4, 0x0f), ('f4', 0, 0x0f)]
 383                         self.assertEqual(len(self.S8P._fields), 2)
 384                         self.assertEqual(calc(self.S8P), om)
 385
 386         def test_to_bytes(self):
 387                 with self.subTest('16 bit total (MSB): f4a + f8 + f4b'):
 388                         vals = { 'f4a' : 0x0f, 'f8' : 0xff, 'f4b' : 0x0f }
 389                         self.assertEqual(self.S16.to_bytes(vals), b'\xff\xff')
 390                         vals = { 'f4a' : 0x00, 'f8' : 0x00, 'f4b' : 0x00 }
 391                         self.assertEqual(self.S16.to_bytes(vals), b'\x00\x00')
 392                         vals = { 'f4a' : 0x0f, 'f8' : 0x00, 'f4b' : 0x0f }
 393                         self.assertEqual(self.S16.to_bytes(vals), b'\xf0\x0f')
 394                         vals = { 'f4a' : 0x00, 'f8' : 0xff, 'f4b' : 0x00 }
 395                         self.assertEqual(self.S16.to_bytes(vals), b'\x0f\xf0')
 396
 397                 with self.subTest('8 bit total (MSB): f4 + f1 + f3'):
 398                         vals = { 'f4' : 0x0f, 'f1' : 0x01, 'f3' : 0x07 }
 399                         self.assertEqual(self.S8M.to_bytes(vals), b'\xff')
 400                         vals = { 'f4' : 0x00, 'f1' : 0x00, 'f3' : 0x00 }
 401                         self.assertEqual(self.S8M.to_bytes(vals), b'\x00')
 402                         vals = { 'f4' : 0x0f, 'f1' : 0x00, 'f3' : 0x00 }
 403                         self.assertEqual(self.S8M.to_bytes(vals), b'\xf0')
 404
 405                 with self.subTest('8 bit total (LSB): f4 + f1 + f3'):
 406                         vals = { 'f4' : 0x0f, 'f1' : 0x01, 'f3' : 0x07 }
 407                         self.assertEqual(self.S8L.to_bytes(vals), b'\xff')
 408                         vals = { 'f4' : 0x00, 'f1' : 0x00, 'f3' : 0x00 }
 409                         self.assertEqual(self.S8L.to_bytes(vals), b'\x00')
 410                         vals = { 'f4' : 0x0f, 'f1' : 0x00, 'f3' : 0x00 }
 411                         self.assertEqual(self.S8L.to_bytes(vals), b'\x0f')
 412
 413         def test_from_bytes(self):
 414                 pad = b'\xff' * 64
 415
 416                 with self.subTest('16 bit total (MSB): f4a + f8 + f4b'):
 417                         vals = { 'f4a' : 0x0f, 'f8' : 0xff, 'f4b' : 0x0f }
 418                         self.assertEqual(self.from_bytes(self.S16, b'\xff\xff' + pad), vals)
 419                         vals = { 'f4a' : 0x00, 'f8' : 0x00, 'f4b' : 0x00 }
 420                         self.assertEqual(self.from_bytes(self.S16, b'\x00\x00' + pad), vals)
 421                         vals = { 'f4a' : 0x0f, 'f8' : 0x00, 'f4b' : 0x0f }
 422                         self.assertEqual(self.from_bytes(self.S16, b'\xf0\x0f' + pad), vals)
 423                         vals = { 'f4a' : 0x00, 'f8' : 0xff, 'f4b' : 0x00 }
 424                         self.assertEqual(self.from_bytes(self.S16, b'\x0f\xf0' + pad), vals)
 425
 426                 with self.subTest('8 bit total (MSB): f4 + f1 + f3'):
 427                         vals = { 'f4' : 0x0f, 'f1' : 0x01, 'f3' : 0x07 }
 428                         self.assertEqual(self.from_bytes(self.S8M, b'\xff' + pad), vals)
 429                         vals = { 'f4' : 0x00, 'f1' : 0x00, 'f3' : 0x00 }
 430                         self.assertEqual(self.from_bytes(self.S8M, b'\x00' + pad), vals)
 431                         vals = { 'f4' : 0x0f, 'f1' : 0x00, 'f3' : 0x00 }
 432                         self.assertEqual(self.from_bytes(self.S8M, b'\xf0' + pad), vals)
 433
 434                 with self.subTest('8 bit total (LSB): f4 + f1 + f3'):
 435                         vals = { 'f4' : 0x0f, 'f1' : 0x01, 'f3' : 0x07 }
 436                         self.assertEqual(self.from_bytes(self.S8L, b'\xff' + pad), vals)
 437                         vals = { 'f4' : 0x00, 'f1' : 0x00, 'f3' : 0x00 }
 438                         self.assertEqual(self.from_bytes(self.S8L, b'\x00' + pad), vals)
 439                         vals = { 'f4' : 0x0f, 'f1' : 0x00, 'f3' : 0x00 }
 440                         self.assertEqual(self.from_bytes(self.S8L, b'\x0f' + pad), vals)
 441
 442         def test_to_bytes_val(self):
 443                 with self.subTest('fixed values in absence of user-supplied values'):
 444                         vals = { 'f3' : 0x00 } # | { 'f4' : 2, 'f1' : 0 }
 445                         self.assertEqual(self.S8V.to_bytes(vals), b'\x20')
 446
 447                 with self.subTest('fixed values take precedence'):
 448                         vals = { 'f4' : 1, 'f1' : 1, 'f3' : 0 }
 449                         self.assertEqual(self.S8V.to_bytes(vals), b'\x20')
 450
 451         def test_from_bytes_val(self):
 452                 with self.assertRaises(codec.DecodeError):
 453                         self.S8V.from_bytes({ }, b'\xf0') # 'f4': 15 vs 2
 454
 455                 with self.assertRaises(codec.DecodeError):
 456                         self.S8V.from_bytes({ }, b'\x08') # 'f1': 1 vs 0
 457
 458                 # Field 'f3' takes any value, no exceptions shall be raised
 459                 for i in range(8):
 460                         data, vals = bytes([0x20 + i]), { 'f4' : 2, 'f1' : 0, 'f3' : i }
 461                         self.assertEqual(self.from_bytes(self.S8V, data), vals)
 462
 463         def test_to_bytes_spare(self):
 464                 self.assertEqual(self.S8P.to_bytes({ 'f4' : 0x00 }), b'\x00')
 465                 self.assertEqual(self.S8P.to_bytes({ 'f4' : 0x0f }), b'\x0f')
 466                 self.assertEqual(self.S8P.to_bytes({ 'f4' : 0xff }), b'\x0f')
 467
 468         def test_from_bytes_spare(self):
 469                 self.assertEqual(self.from_bytes(self.S8P, b'\x00'), { 'f4' : 0x00 })
 470                 self.assertEqual(self.from_bytes(self.S8P, b'\x0f'), { 'f4' : 0x0f })
 471                 self.assertEqual(self.from_bytes(self.S8P, b'\xff'), { 'f4' : 0x0f })
 472
 473 class TestPDU(codec.Envelope):
 474         STRUCT = (
 475                 codec.BitFieldSet(len=2, set=(
 476                         codec.BitField('ver', bl=4),
 477                         codec.BitField('flag', bl=1),
 478                 )),
 479                 codec.Uint16BE('len'),
 480                 codec.Buf('data'),
 481                 codec.Buf('tail', len=2),
 482         )
 483
 484         def __init__(self, *args, **kw):
 485                 codec.Envelope.__init__(self, *args, **kw)
 486                 self.STRUCT[-3].get_val = lambda v: len(v['data'])
 487                 self.STRUCT[-2].get_len = lambda v, _: v['len']
 488                 self.STRUCT[-1].get_pres = lambda v: bool(v['flag'])
 489
 490         def check(self, vals: dict) -> None:
 491                 if not vals['ver'] in (0, 1, 2):
 492                         raise ValueError('Unknown version %d' % vals['ver'])
 493
 494 class Envelope(unittest.TestCase):
 495         def test_rest_octets(self):
 496                 pdu = TestPDU(check_len=False)
 497                 pdu.from_bytes(b'\x00' * 64)
 498
 499                 with self.assertRaises(codec.DecodeError):
 500                         pdu = TestPDU(check_len=True)
 501                         pdu.from_bytes(b'\x00' * 64) # 'len' : 0
 502
 503         def test_field_raises(self):
 504                 pdu = TestPDU()
 505                 with self.assertRaises(codec.EncodeError):
 506                         pdu.c = { 'ver' : 0, 'flag' : 1, 'data' : b'\xff' * 16 }
 507                         pdu.to_bytes() # KeyError: 'tail' not found
 508
 509         def test_to_bytes(self):
 510                 pdu = TestPDU()
 511
 512                 # No content in the new instances
 513                 self.assertEqual(pdu.c, { })
 514
 515                 pdu.c = { 'ver' : 0, 'flag' : 1, 'data' : b'', 'tail' : b'\xde\xbe' }
 516                 self.assertEqual(pdu.to_bytes(), b'\x08\x00\x00\x00' + b'\xde\xbe')
 517
 518                 pdu.c = { 'ver' : 1, 'flag' : 0, 'data' : b'\xff' * 15 }
 519                 self.assertEqual(pdu.to_bytes(), b'\x10\x00\x00\x0f' + b'\xff' * 15)
 520
 521                 pdu.c = { 'ver' : 2, 'flag' : 1, 'data' : b'\xf0', 'tail' : b'\xbe\xed' }
 522                 self.assertEqual(pdu.to_bytes(), b'\x28\x00\x00\x01\xf0\xbe\xed')
 523
 524         def test_from_bytes(self):
 525                 pdu = TestPDU()
 526
 527                 # No content in the new instances
 528                 self.assertEqual(pdu.c, { })
 529
 530                 c = { 'ver' : 0, 'flag' : 1, 'len' : 0, 'data' : b'', 'tail' : b'\xde\xbe' }
 531                 pdu.from_bytes(b'\x08\x00\x00\x00' + b'\xde\xbe')
 532                 self.assertEqual(pdu.c, c)
 533
 534                 c = { 'ver' : 1, 'flag' : 0, 'len' : 15, 'data' : b'\xff' * 15 }
 535                 pdu.from_bytes(b'\x10\x00\x00\x0f' + b'\xff' * 15)
 536                 self.assertEqual(pdu.c, c)
 537
 538                 c = { 'ver' : 2, 'flag' : 1, 'len' : 1, 'data' : b'\xf0', 'tail' : b'\xbe\xed' }
 539                 pdu.from_bytes(b'\x28\x00\x00\x01\xf0\xbe\xed')
 540                 self.assertEqual(pdu.c, c)
 541
 542         def test_to_bytes_check(self):
 543                 pdu = TestPDU()
 544
 545                 pdu.c = { 'ver' : 8, 'flag' : 1, 'data' : b'', 'tail' : b'\xde\xbe' }
 546                 with self.assertRaises(ValueError):
 547                         pdu.to_bytes()
 548
 549         def test_from_bytes_check(self):
 550                 pdu = TestPDU()
 551
 552                 with self.assertRaises(ValueError):
 553                         pdu.from_bytes(b'\xf0\x00\x00\x00')
 554
 555 class Sequence(unittest.TestCase):
 556         class TLV(codec.Envelope):
 557                 STRUCT = (
 558                         codec.Uint('T'),
 559                         codec.Uint('L'),
 560                         codec.Buf('V'),
 561                 )
 562
 563                 def __init__(self, *args, **kw) -> None:
 564                         codec.Envelope.__init__(self, *args, **kw)
 565                         self.STRUCT[-2].get_val = lambda v: len(v['V'])
 566                         self.STRUCT[-1].get_len = lambda v, _: v['L']
 567
 568         # Sequence of TLVs
 569         SEQ = codec.Sequence(item=TLV())
 570
 571         Vseq, Bseq = [
 572                 { 'T' : 0xde, 'L' : 4, 'V' : b'\xde\xad\xbe\xef' },
 573                 { 'T' : 0xbe, 'L' : 2, 'V' : b'\xbe\xef' },
 574                 { 'T' : 0xbe, 'L' : 2, 'V' : b'\xef\xbe' },
 575                 { 'T' : 0x00, 'L' : 0, 'V' : b'' },
 576         ], b''.join([
 577                 b'\xde\x04\xde\xad\xbe\xef',
 578                 b'\xbe\x02\xbe\xef',
 579                 b'\xbe\x02\xef\xbe',
 580                 b'\x00\x00',
 581         ])
 582
 583         def test_to_bytes(self):
 584                 res = self.SEQ.to_bytes(self.Vseq)
 585                 self.assertEqual(res, self.Bseq)
 586
 587         def test_from_bytes(self):
 588                 res = self.SEQ.from_bytes(self.Bseq)
 589                 self.assertEqual(res, self.Vseq)
 590
 591 if __name__ == '__main__':
 592         unittest.main()