Test encoding/decoding of a message and records into 'RecordIO' format.
()
| 51 | |
| 52 | |
| 53 | def test_encode_decode(): |
| 54 | """ |
| 55 | Test encoding/decoding of a message and records into 'RecordIO' format. |
| 56 | """ |
| 57 | total_messages = 10 |
| 58 | |
| 59 | try: |
| 60 | encoder = recordio.Encoder(lambda s: bytes(json.dumps(s), "UTF-8")) |
| 61 | except Exception as exception: |
| 62 | raise MesosException("Error instantiating 'RecordIO' encoder: {error}" |
| 63 | .format(error=exception)) |
| 64 | |
| 65 | try: |
| 66 | decoder = recordio.Decoder(lambda s: json.loads(s.decode("UTF-8"))) |
| 67 | except Exception as exception: |
| 68 | raise MesosException("Error instantiating 'RecordIO' decoder: {error}" |
| 69 | .format(error=exception)) |
| 70 | |
| 71 | try: |
| 72 | message = { |
| 73 | "type": "ATTACH_CONTAINER_OUTPUT", |
| 74 | "containerId": "123456789" |
| 75 | } |
| 76 | |
| 77 | encoded = b"" |
| 78 | for _ in range(total_messages): |
| 79 | encoded += encoder.encode(message) |
| 80 | |
| 81 | except Exception as exception: |
| 82 | raise MesosException("Error encoding 'RecordIO' message: {error}" |
| 83 | .format(error=exception)) |
| 84 | |
| 85 | try: |
| 86 | all_records = [] |
| 87 | offset = 0 |
| 88 | chunk_size = 5 |
| 89 | while offset < len(encoded): |
| 90 | records = decoder.decode(encoded[offset:offset + chunk_size]) |
| 91 | all_records.extend(records) |
| 92 | offset += chunk_size |
| 93 | |
| 94 | assert len(all_records) == total_messages |
| 95 | |
| 96 | for record in all_records: |
| 97 | assert record == message |
| 98 | except Exception as exception: |
| 99 | raise MesosException("Error decoding 'RecordIO' messages: {error}" |
| 100 | .format(error=exception)) |