def test_upload_is_only_valid_if_chunk_count_matches_expected(self):
        storage = LongMessageStorage(MemoryStorage(), MemoryStorage())
        chunks = [b'12345', b'1234568', b'98765432']

        checksum = hashlib.md5()
        for chunk in chunks:
            checksum.update(chunk)

        expected_md5 = checksum.hexdigest()

        handler = LongMessageHandler(storage)
        handler.select_long_message_type(LongMessageType.FIRMWARE_DATA)

        expectations = ((len(chunks), LongMessageStatus.READY),
                        (0, LongMessageStatus.READY),
                        (len(chunks) - 1, LongMessageStatus.VALIDATION_ERROR),
                        (len(chunks) + 1, LongMessageStatus.VALIDATION_ERROR))

        for length, expected_status in expectations:
            handler.init_transfer(expected_md5, length)
            for chunk in chunks:
                handler.upload_message(chunk)
            handler.finalize_message()

            status = handler.read_status()
            self.assertEqual(expected_status, status.status)
    def test_finalize_validates_and_stores_uploaded_message(self, mock_hash):
        storage = Mock()
        storage.read_status = Mock(return_value=LongMessageStatusInfo(
            LongMessageStatus.READY, 'new_md5', 123))
        storage.set_long_message = Mock()

        mock_hash.return_value = mock_hash
        mock_hash.update = Mock()
        mock_hash.hexdigest = Mock()

        handler = LongMessageHandler(storage)

        # message invalid
        for mt in self.known_message_types:
            mock_hash.reset_mock()

            mock_hash.hexdigest.return_value = 'invalid_md5'
            with self.subTest(mt='invalid ({})'.format(mt)):
                handler.select_long_message_type(mt)
                handler.init_transfer('new_md5')
                self.assertEqual(1, mock_hash.call_count)

                handler.upload_message(b'12345')
                self.assertEqual(1, mock_hash.update.call_count)
                handler.upload_message(b'67890')
                self.assertEqual(2, mock_hash.update.call_count)

                handler.finalize_message()
                self.assertEqual(1, mock_hash.hexdigest.call_count)

                status = handler.read_status()
                self.assertEqual(LongMessageStatus.VALIDATION_ERROR,
                                 status.status)
                self.assertEqual(0, storage.set_long_message.call_count)

        # message valid
        for mt in self.known_message_types:
            mock_hash.reset_mock()
            storage.reset_mock()

            mock_hash.hexdigest.return_value = 'new_md5'
            with self.subTest(mt='valid ({})'.format(mt)):
                handler.select_long_message_type(mt)
                handler.init_transfer('new_md5')
                self.assertEqual(1, mock_hash.call_count)

                handler.upload_message(b'12345')
                self.assertEqual(1, mock_hash.update.call_count)
                handler.upload_message(b'67890')
                self.assertEqual(2, mock_hash.update.call_count)

                handler.finalize_message()
                self.assertEqual(1, mock_hash.hexdigest.call_count)
                self.assertEqual(1, storage.set_long_message.call_count)

                status = handler.read_status()
                self.assertEqual(1, storage.read_status.call_count)
                self.assertEqual(LongMessageStatus.READY, status.status)
    def test_reading_during_upload_returns_md5_and_current_length_of_new_message(
            self):
        storage = Mock()

        handler = LongMessageHandler(storage)

        for mt in self.known_message_types:
            with self.subTest(mt=mt):
                handler.select_long_message_type(mt)
                handler.init_transfer('new_md5')
                handler.upload_message(b'12345')
                status = handler.read_status()
                self.assertEqual(LongMessageStatus.UPLOAD, status.status)
                self.assertEqual('new_md5', status.md5)
                self.assertEqual(5, status.length)

                handler.upload_message(b'67890')
                status = handler.read_status()
                self.assertEqual(LongMessageStatus.UPLOAD, status.status)
                self.assertEqual('new_md5', status.md5)
                self.assertEqual(10, status.length)
    def test_finalize_notifies_about_valid_message(self, mock_hash):
        storage = Mock()
        storage.read_status = Mock(return_value=LongMessageStatusInfo(
            LongMessageStatus.READY, 'new_md5', 123))

        mock_hash.return_value = mock_hash
        mock_hash.update = Mock()
        mock_hash.hexdigest = Mock()

        mock_callback = Mock()

        handler = LongMessageHandler(storage)
        handler.on_message_updated(mock_callback)

        # message invalid
        for mt in self.known_message_types:
            mock_hash.reset_mock()

            mock_hash.hexdigest.return_value = 'invalid_md5'
            with self.subTest(mt='invalid ({})'.format(mt)):
                handler.select_long_message_type(mt)
                handler.init_transfer('new_md5')
                handler.upload_message(b'12345')
                handler.finalize_message()

                self.assertEqual(0, mock_callback.call_count)

        # message valid
        for mt in self.known_message_types:
            mock_callback.reset_mock()

            mock_hash.hexdigest.return_value = 'new_md5'
            with self.subTest(mt='valid ({})'.format(mt)):
                handler.select_long_message_type(mt)
                handler.init_transfer('new_md5')
                handler.upload_message(b'12345')
                handler.finalize_message()

                self.assertEqual(1, mock_callback.call_count)
    def test_init_is_required_before_write_and_finalize(self):
        handler = LongMessageHandler(None)

        self.assertRaises(LongMessageError,
                          lambda: handler.upload_message([1]))
        self.assertRaises(LongMessageError, handler.finalize_message)