Commit b9e2d3b2 authored by Dom Sekotill's avatar Dom Sekotill
Browse files

Add abort_on_unknown option to FilterProtocol 

This option allows users to pretend that Abort was sent when an unknown
message is received.
parent ecf12c1a

kilter/protocol/core.py

+5 −2
Original line number Diff line number Diff line
@@ -210,7 +210,8 @@ class FilterProtocol: 
	buffers.  The class checks the correctness of responses sent back to the MTA.

	"""



	def __init__(self) -> None:

	def __init__(self, *, abort_on_unknown: bool = False) -> None:

		self.abort_on_unknown = abort_on_unknown

		self.nr = set[bytes]()

		self.actions = set[bytes]([messages.Progress.ident])

		self.state: tuple[messages.Message, set[bytes]]|None = None
@@ -236,7 +237,9 @@ class FilterProtocol: 
				return

			except UnknownMessage as exc:

				del buf[:len(exc.contents)]

				if not self.abort_on_unknown:

					raise

				yield Abort()

			else:

				yield self._check_recv(message)

				message.release()

tests/test_core_filter_protocol.py

+13 −0
Original line number Diff line number Diff line
@@ -76,6 +76,19 @@ class FilterProtocolTests(unittest.TestCase): 


		assert exc_cm.exception.contents == msg



	def test_read_unimplemented_abort(self) -> None:

		"""

		Check that unknown messages cause Abort to be returned when enabled

		"""

		buf = SimpleBuffer(20)

		buf[:] = b"\x00\x00\x00\x01S"



		for msg in FilterProtocol(abort_on_unknown=True).read_from(buf):

			assert isinstance(msg, Abort)

			break

		else:

			self.fail("No messages read")



	def test_read_unexpected(self) -> None:

		"""

		Check that reading an available message before a response raises UnexpectedMessage