Repository: duesee/imap-codec (language: rust)

Issue:
bindings(Python): Use Python classes for `Fragment` variants 
Currently, the fragments of `Encoded` (Rust enum variants) are deserializes into variant dictionaries, e.g. `{"Line": {"data": list(b"a NOOP\r\n")}}`.

By using individual classes for each variant (`LineFragment`, `LiteralFragment`), handling these in the Python code would be greatly improved.

Related: This would remove the need for `Fragment` to implement `Serialize`, see #555.


Discussion / hints:

Fix patch:
diff --git a/bindings/imap-codec-python/imap_codec.pyi b/bindings/imap-codec-python/imap_codec.pyi
index 89874e45..33e0544a 100644
--- a/bindings/imap-codec-python/imap_codec.pyi
+++ b/bindings/imap-codec-python/imap_codec.pyi
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Tuple
+from typing import Tuple, Union
 
 class DecodeError(Exception):
     """
@@ -25,6 +25,72 @@ class DecodeLiteralFound(DecodeError):
     The decoder stopped at the beginning of literal data.
     """
 
+class LiteralMode:
+    """
+    Literal mode, i.e., sync or non-sync.
+
+    - Sync: A synchronizing literal, i.e., `{<n>}\r\n<data>`.
+    - NonSync: A non-synchronizing literal according to RFC 7888, i.e., `{<n>+}\r\n<data>`.
+
+    Warning: The non-sync literal extension must only be used when the server advertised support
+             for it sending the LITERAL+ or LITERAL- capability.
+    """
+
+    Sync: LiteralMode
+    NonSync: LiteralMode
+
+class LineFragment:
+    """
+    Fragment of a line that is ready to be send.
+    """
+
+    def __init__(self, data: bytes):
+        """
+        Create a line fragment from data bytes
+
+        :param data: Data bytes of fragment
+        :raises TypeError: `data` is not byte-like
+        """
+
+    @property
+    def data(self) -> bytes:
+        """
+        Get line fragment data bytes
+
+        :return: Data bytes of fragment
+        """
+
+class LiteralFragment:
+    """
+    Fragment of a literal that may require an action before it should be send.
+    """
+
+    def __init__(self, data: bytes, mode: LiteralMode):
+        """
+        Create a literal fragment from data bytes and literal mode
+
+        :param data: Data bytes of fragment
+        :param mode: Literal mode
+        :raises TypeError: `data` is not byte-like
+        :raises TypeError: `mode` is invalid
+        """
+
+    @property
+    def data(self) -> bytes:
+        """
+        Get literal fragment data bytes
+
+        :return: Data bytes of fragment
+        """
+
+    @property
+    def mode(self) -> LiteralMode:
+        """
+        Get literal fragment literal mode
+
+        :return: Literal mode
+        """
+
 class Encoded:
     """
     An encoded message.
@@ -36,7 +102,7 @@ class Encoded:
     """
 
     def __iter__(self) -> Encoded: ...
-    def __next__(self) -> dict: ...
+    def __next__(self) -> Union[LineFragment, LiteralFragment]: ...
     def dump(self) -> bytes:
         """
         Dump the (remaining) encoded data without being guided by fragments.
diff --git a/bindings/imap-codec-python/src/encoded.rs b/bindings/imap-codec-python/src/encoded.rs
new file mode 100644
index 00000000..d896933e
--- /dev/null
+++ b/bindings/imap-codec-python/src/encoded.rs
@@ -0,0 +1,150 @@
+use imap_codec::{
+    encode::{Encoded, Fragment},
+    imap_types::core::LiteralMode,
+};
+use pyo3::{prelude::*, types::PyBytes};
+
+/// Python class representing a literal mode
+#[derive(Debug, Clone, Copy, PartialEq)]
+#[pyclass(name = "LiteralMode", eq)]
+pub(crate) enum PyLiteralMode {
+    Sync,
+    NonSync,
+}
+
+impl std::fmt::Display for PyLiteralMode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PyLiteralMode::Sync => f.write_str("LiteralMode.Sync"),
+            PyLiteralMode::NonSync => f.write_str("LiteralMode.NonSync"),
+        }
+    }
+}
+
+impl From<LiteralMode> for PyLiteralMode {
+    fn from(value: LiteralMode) -> Self {
+        match value {
+            LiteralMode::Sync => PyLiteralMode::Sync,
+            LiteralMode::NonSync => PyLiteralMode::NonSync,
+        }
+    }
+}
+
+/// Python class representing a line fragment
+#[derive(Debug, Clone, PartialEq)]
+#[pyclass(name = "LineFragment", eq)]
+pub(crate) struct PyLineFragment {
+    data: Vec<u8>,
+}
+
+#[pymethods]
+impl PyLineFragment {
+    /// Create a new line fragment from data
+    ///
+    /// `data` can be anything that can be extracted to `Vec`, e.g. Python `bytes`
+    #[new]
+    pub(crate) fn new(data: Vec<u8>) -> Self {
+        Self { data }
+    }
+
+    /// Retrieve the data from the fragment as Python `bytes`
+    #[getter]
+    pub(crate) fn data<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
+        PyBytes::new_bound(py, self.data.as_slice())
+    }
+
+    /// String representation of the fragment, e.g. `b'hello'`
+    pub(crate) fn __str__(&self, py: Python) -> String {
+        self.data(py).to_string()
+    }
+
+    /// Printable representation of the fragment, e.g. `LineFragment(b'hello')`
+    pub(crate) fn __repr__(&self, py: Python) -> String {
+        format!("LineFragment({})", self.__str__(py))
+    }
+}
+
+/// Python class representing a literal fragment
+#[derive(Debug, Clone, PartialEq)]
+#[pyclass(name = "LiteralFragment", eq)]
+pub(crate) struct PyLiteralFragment {
+    data: Vec<u8>,
+    mode: PyLiteralMode,
+}
+
+#[pymethods]
+impl PyLiteralFragment {
+    /// Create a new literal fragment from data and mode
+    ///
+    /// `data` can be anything that can be extracted to `Vec`, e.g. Python `bytes`
+    #[new]
+    pub(crate) fn try_new(data: Vec<u8>, mode: PyLiteralMode) -> PyResult<Self> {
+        Ok(Self { data, mode })
+    }
+
+    /// Retrieve the data of the fragment as Python `bytes`
+    #[getter]
+    pub(crate) fn data<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
+        PyBytes::new_bound(py, self.data.as_slice())
+    }
+
+    /// Retrieve the mode of the fragment
+    #[getter]
+    pub(crate) fn mode(&self) -> PyLiteralMode {
+        self.mode
+    }
+
+    /// String representation of the fragment, e.g. `(b'hello', 'Sync')`
+    pub(crate) fn __str__(&self, py: Python) -> String {
+        format!("({}, {})", self.data(py), self.mode)
+    }
+
+    /// Printable representation of the fragment, e.g. `LiteralFragment(b'hello', 'Sync')`
+    pub(crate) fn __repr__(&self, py: Python) -> String {
+        format!("LiteralFragment{}", self.__str__(py))
+    }
+}
+
+/// Python wrapper classes for `Encoded`
+///
+/// This implements a Python iterator over the containing fragments.
+#[derive(Debug, Clone)]
+#[pyclass(name = "Encoded")]
+pub(crate) struct PyEncoded(pub(crate) Option<Encoded>);
+
+#[pymethods]
+impl PyEncoded {
+    /// Initialize iterator
+    pub(crate) fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    /// Return next fragment
+    pub(crate) fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<PyObject>> {
+        // Try to get next `Fragment` from `Encoded` iterator
+        let Some(fragment) = slf.0.as_mut().and_then(|encoded| encoded.next()) else {
+            return Ok(None);
+        };
+
+        // Return instance of `PyLineFragment` or `PyLiteralFragment` as a generic `PyObject`.
+        Ok(Some(match fragment {
+            Fragment::Line { data } => {
+                Bound::new(slf.py(), PyLineFragment::new(data))?.to_object(slf.py())
+            }
+            Fragment::Literal { data, mode } => {
+                Bound::new(slf.py(), PyLiteralFragment::try_new(data, mode.into())?)?
+                    .to_object(slf.py())
+            }
+        }))
+    }
+
+    /// Dump remaining fragment data
+    pub(crate) fn dump(mut slf: PyRefMut<'_, Self>) -> PyResult<Bound<PyBytes>> {
+        let encoded = slf.0.take();
+        let dump = match encoded {
+            Some(encoded) => encoded.dump(),
+            None => Vec::new(),
+        };
+        Ok(PyBytes::new_bound(slf.py(), &dump))
+    }
+}
diff --git a/bindings/imap-codec-python/src/lib.rs b/bindings/imap-codec-python/src/lib.rs
index 23ee0717..da568fb4 100644
--- a/bindings/imap-codec-python/src/lib.rs
+++ b/bindings/imap-codec-python/src/lib.rs
@@ -1,6 +1,9 @@
+mod encoded;
+
+use encoded::PyEncoded;
 use imap_codec::{
     decode::{self, Decoder},
-    encode::{Encoded, Encoder},
+    encode::Encoder,
     AuthenticateDataCodec, CommandCodec, GreetingCodec, IdleDoneCodec, ResponseCodec,
 };
 use pyo3::{create_exception, exceptions::PyException, prelude::*, types::PyBytes};
@@ -11,46 +14,10 @@ create_exception!(imap_codec, DecodeFailed, DecodeError);
 create_exception!(imap_codec, DecodeIncomplete, DecodeError);
 create_exception!(imap_codec, DecodeLiteralFound, DecodeError);
 
-/// Wrapper for `Encoded`
-///
-/// This implements a Python iterator over the containing fragments.
-#[derive(Debug, Clone)]
-#[pyclass(name = "Encoded")]
-struct PyEncoded(Option<Encoded>);
-
-#[pymethods]
-impl PyEncoded {
-    /// Initialize iterator
-    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
-        slf
-    }
-
-    /// Return next fragment
-    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<Bound<PyAny>>> {
-        let Some(encoded) = &mut slf.0 else {
-            return Ok(None);
-        };
-        Ok(encoded
-            .next()
-            .map(|value| serde_pyobject::to_pyobject(slf.py(), &value))
-            .transpose()?)
-    }
-
-    /// Dump remaining fragment data
-    fn dump(mut slf: PyRefMut<'_, Self>) -> PyResult<Bound<PyBytes>> {
-        let encoded = slf.0.take();
-        let dump = match encoded {
-            Some(encoded) => encoded.dump(),
-            None => Vec::new(),
-        };
-        Ok(PyBytes::new_bound(slf.py(), &dump))
-    }
-}
-
-/// Wrapper for `GreetingCodec`
+/// Python class for using `GreetingCodec`
 #[derive(Debug, Clone, PartialEq)]
 #[pyclass(name = "GreetingCodec")]
-struct PyGreetingCodec(GreetingCodec);
+struct PyGreetingCodec;
 
 #[pymethods]
 impl PyGreetingCodec {
@@ -80,10 +47,10 @@ impl PyGreetingCodec {
     }
 }
 
-/// Wrapper for `CommandCodec`
+/// Python class for using `CommandCodec`
 #[derive(Debug, Clone, PartialEq)]
 #[pyclass(name = "CommandCodec")]
-struct PyCommandCodec(CommandCodec);
+struct PyCommandCodec;
 
 #[pymethods]
 impl PyCommandCodec {
@@ -119,10 +86,10 @@ impl PyCommandCodec {
     }
 }
 
-/// Wrapper for `AuthenticateDataCodec`
+/// Python class for using `AuthenticateDataCodec`
 #[derive(Debug, Clone, PartialEq)]
 #[pyclass(name = "AuthenticateDataCodec")]
-struct PyAuthenticateDataCodec(AuthenticateDataCodec);
+struct PyAuthenticateDataCodec;
 
 #[pymethods]
 impl PyAuthenticateDataCodec {
@@ -151,10 +118,10 @@ impl PyAuthenticateDataCodec {
     }
 }
 
-/// Wrapper for `ResponseCodec`
+/// Python class for using `ResponseCodec`
 #[derive(Debug, Clone, PartialEq)]
 #[pyclass(name = "ResponseCodec")]
-struct PyResponseCodec(ResponseCodec);
+struct PyResponseCodec;
 
 #[pymethods]
 impl PyResponseCodec {
@@ -188,10 +155,10 @@ impl PyResponseCodec {
     }
 }
 
-/// Wrapper for `IdleDoneCodec`
+/// Python class for using `IdleDoneCodec`
 #[derive(Debug, Clone, PartialEq)]
 #[pyclass(name = "IdleDoneCodec")]
-struct PyIdleDoneCodec(IdleDoneCodec);
+struct PyIdleDoneCodec;
 
 #[pymethods]
 impl PyIdleDoneCodec {
@@ -233,6 +200,9 @@ fn imap_codec_python(m: &Bound<'_, PyModule>) -> PyResult<()> {
         "DecodeLiteralFound",
         m.py().get_type_bound::<DecodeLiteralFound>(),
     )?;
+    m.add_class::<encoded::PyLiteralMode>()?;
+    m.add_class::<encoded::PyLineFragment>()?;
+    m.add_class::<encoded::PyLiteralFragment>()?;
     m.add_class::<PyEncoded>()?;
     m.add_class::<PyGreetingCodec>()?;
     m.add_class::<PyCommandCodec>()?;


Test patch:
diff --git a/bindings/imap-codec-python/tests/test_authenticate_data_encode.py b/bindings/imap-codec-python/tests/test_authenticate_data_encode.py
index 5ea8a4d5..d21ff726 100644
--- a/bindings/imap-codec-python/tests/test_authenticate_data_encode.py
+++ b/bindings/imap-codec-python/tests/test_authenticate_data_encode.py
@@ -1,6 +1,6 @@
 import unittest
 
-from imap_codec import AuthenticateDataCodec, Encoded
+from imap_codec import AuthenticateDataCodec, Encoded, LineFragment
 
 
 class TestAuthenticateDataEncode(unittest.TestCase):
@@ -9,7 +9,7 @@ def test_authenticate_data(self):
         encoded = AuthenticateDataCodec.encode(authenticate_data)
         self.assertIsInstance(encoded, Encoded)
         fragments = list(encoded)
-        self.assertEqual(fragments, [{"Line": {"data": list(b"VGVzdA==\r\n")}}])
+        self.assertEqual(fragments, [LineFragment(b"VGVzdA==\r\n")])
 
     def test_authenticate_data_dump(self):
         authenticate_data = {"Continue": list(b"Test")}
diff --git a/bindings/imap-codec-python/tests/test_command_encode.py b/bindings/imap-codec-python/tests/test_command_encode.py
index a61eef37..c51d7e61 100644
--- a/bindings/imap-codec-python/tests/test_command_encode.py
+++ b/bindings/imap-codec-python/tests/test_command_encode.py
@@ -1,6 +1,6 @@
 import unittest
 
-from imap_codec import CommandCodec, Encoded
+from imap_codec import CommandCodec, Encoded, LineFragment, LiteralFragment, LiteralMode
 
 
 class TestCommandEncode(unittest.TestCase):
@@ -9,7 +9,7 @@ def test_simple_command(self):
         encoded = CommandCodec.encode(command)
         self.assertIsInstance(encoded, Encoded)
         fragments = list(encoded)
-        self.assertEqual(fragments, [{"Line": {"data": list(b"a NOOP\r\n")}}])
+        self.assertEqual(fragments, [LineFragment(b"a NOOP\r\n")])
 
     def test_simple_command_dump(self):
         command = {"tag": "a", "body": "Noop"}
@@ -36,9 +36,9 @@ def test_multi_fragment_command(self):
         self.assertEqual(
             fragments,
             [
-                {"Line": {"data": list(b"A LOGIN alice {2}\r\n")}},
-                {"Literal": {"data": list(b"\xCA\xFE"), "mode": "Sync"}},
-                {"Line": {"data": list(b"\r\n")}},
+                LineFragment(b"A LOGIN alice {2}\r\n"),
+                LiteralFragment(b"\xCA\xFE", LiteralMode.Sync),
+                LineFragment(b"\r\n"),
             ],
         )
 
@@ -50,7 +50,5 @@ def test_multi_fragment_command_dump(self):
     def test_multi_fragment_command_dump_remaining(self):
         encoded = CommandCodec.encode(self._MULTI_FRAGMENT_COMMAND)
         self.assertIsInstance(encoded, Encoded)
-        self.assertEqual(
-            next(encoded), {"Line": {"data": list(b"A LOGIN alice {2}\r\n")}}
-        )
+        self.assertEqual(next(encoded), LineFragment(b"A LOGIN alice {2}\r\n"))
         self.assertEqual(encoded.dump(), b"\xCA\xFE\r\n")
diff --git a/bindings/imap-codec-python/tests/test_fragments.py b/bindings/imap-codec-python/tests/test_fragments.py
new file mode 100644
index 00000000..2803048f
--- /dev/null
+++ b/bindings/imap-codec-python/tests/test_fragments.py
@@ -0,0 +1,101 @@
+import unittest
+
+from imap_codec import LineFragment, LiteralFragment, LiteralMode
+
+
+class TestLineFragment(unittest.TestCase):
+    def test_data(self):
+        data = b"a NOOP\r\n"
+        fragment = LineFragment(data)
+        self.assertEqual(fragment.data, data)
+
+    def test_repr(self):
+        data = b"a NOOP\r\n"
+        fragment = LineFragment(data)
+        self.assertEqual(repr(fragment), f"LineFragment({data})")
+
+    def test_str(self):
+        data = b"a NOOP\r\n"
+        fragment = LineFragment(data)
+        self.assertEqual(str(fragment), str(data))
+
+    def test_eq(self):
+        fragment1 = LineFragment(b"a NOOP\r\n")
+        fragment2 = LineFragment(b"a NOOP\r\n")
+        fragment3 = LineFragment(b"a LOGIN alice pass\r\n")
+
+        self.assertEqual(fragment1, fragment1)
+        self.assertEqual(fragment1, fragment2)
+        self.assertNotEqual(fragment1, fragment3)
+
+        self.assertEqual(fragment2, fragment1)
+        self.assertEqual(fragment2, fragment2)
+        self.assertNotEqual(fragment2, fragment3)
+
+        self.assertNotEqual(fragment3, fragment1)
+        self.assertNotEqual(fragment3, fragment2)
+        self.assertEqual(fragment3, fragment3)
+
+
+class TestLiteralFragment(unittest.TestCase):
+    def test_data(self):
+        data = b"\x01\x02\x03\x04"
+        fragment = LiteralFragment(data, LiteralMode.Sync)
+        self.assertEqual(fragment.data, data)
+
+    def test_mode(self):
+        mode = LiteralMode.Sync
+        fragment = LiteralFragment(b"\x01\x02\x03\x04", mode)
+        self.assertEqual(fragment.mode, mode)
+
+        mode = LiteralMode.NonSync
+        fragment = LiteralFragment(b"\x01\x02\x03\x04", mode)
+        self.assertEqual(fragment.mode, mode)
+
+    def test_repr(self):
+        data = b"\x01\x02\x03\x04"
+
+        mode = LiteralMode.Sync
+        fragment = LiteralFragment(data, mode)
+        self.assertEqual(repr(fragment), f"LiteralFragment({data}, {mode})")
+
+        mode = LiteralMode.NonSync
+        fragment = LiteralFragment(data, mode)
+        self.assertEqual(repr(fragment), f"LiteralFragment({data}, {mode})")
+
+    def test_str(self):
+        data = b"\x01\x02\x03\x04"
+
+        mode = LiteralMode.Sync
+        fragment = LiteralFragment(data, mode)
+        self.assertEqual(str(fragment), f"({data}, {mode})")
+
+        mode = LiteralMode.NonSync
+        fragment = LiteralFragment(data, mode)
+        self.assertEqual(str(fragment), f"({data}, {mode})")
+
+    def test_eq(self):
+        fragment1 = LiteralFragment(b"data", LiteralMode.Sync)
+        fragment2 = LiteralFragment(b"data", LiteralMode.Sync)
+        fragment3 = LiteralFragment(b"data", LiteralMode.NonSync)
+        fragment4 = LiteralFragment(b"\x01\x02\x03\x04", LiteralMode.NonSync)
+
+        self.assertEqual(fragment1, fragment1)
+        self.assertEqual(fragment1, fragment2)
+        self.assertNotEqual(fragment1, fragment3)
+        self.assertNotEqual(fragment1, fragment4)
+
+        self.assertEqual(fragment2, fragment1)
+        self.assertEqual(fragment2, fragment2)
+        self.assertNotEqual(fragment2, fragment3)
+        self.assertNotEqual(fragment2, fragment4)
+
+        self.assertNotEqual(fragment3, fragment1)
+        self.assertNotEqual(fragment3, fragment2)
+        self.assertEqual(fragment3, fragment3)
+        self.assertNotEqual(fragment3, fragment4)
+
+        self.assertNotEqual(fragment4, fragment1)
+        self.assertNotEqual(fragment4, fragment2)
+        self.assertNotEqual(fragment4, fragment3)
+        self.assertEqual(fragment4, fragment4)
diff --git a/bindings/imap-codec-python/tests/test_greeting_encode.py b/bindings/imap-codec-python/tests/test_greeting_encode.py
index e9b2077d..b64716fb 100644
--- a/bindings/imap-codec-python/tests/test_greeting_encode.py
+++ b/bindings/imap-codec-python/tests/test_greeting_encode.py
@@ -1,6 +1,6 @@
 import unittest
 
-from imap_codec import Encoded, GreetingCodec
+from imap_codec import Encoded, GreetingCodec, LineFragment
 
 
 class TestGreetingEncode(unittest.TestCase):
@@ -9,9 +9,7 @@ def test_simple_greeting(self):
         encoded = GreetingCodec.encode(greeting)
         self.assertIsInstance(encoded, Encoded)
         fragments = list(encoded)
-        self.assertEqual(
-            fragments, [{"Line": {"data": list(b"* OK Hello, World!\r\n")}}]
-        )
+        self.assertEqual(fragments, [LineFragment(b"* OK Hello, World!\r\n")])
 
     def test_simple_greeting_dump(self):
         greeting = {"code": None, "kind": "Ok", "text": "Hello, World!"}
diff --git a/bindings/imap-codec-python/tests/test_idle_done_encode.py b/bindings/imap-codec-python/tests/test_idle_done_encode.py
index 5208d89a..2e15ced3 100644
--- a/bindings/imap-codec-python/tests/test_idle_done_encode.py
+++ b/bindings/imap-codec-python/tests/test_idle_done_encode.py
@@ -1,6 +1,6 @@
 import unittest
 
-from imap_codec import Encoded, IdleDoneCodec
+from imap_codec import Encoded, IdleDoneCodec, LineFragment
 
 
 class TestIdleDoneEncode(unittest.TestCase):
@@ -9,7 +9,7 @@ def test_idle_done(self):
         encoded = IdleDoneCodec.encode(idle_done)
         self.assertIsInstance(encoded, Encoded)
         fragments = list(encoded)
-        self.assertEqual(fragments, [{"Line": {"data": list(b"DONE\r\n")}}])
+        self.assertEqual(fragments, [LineFragment(b"DONE\r\n")])
 
     def test_idle_done_dump(self):
         idle_done = ()
diff --git a/bindings/imap-codec-python/tests/test_response_encode.py b/bindings/imap-codec-python/tests/test_response_encode.py
index 94149197..7e30c141 100644
--- a/bindings/imap-codec-python/tests/test_response_encode.py
+++ b/bindings/imap-codec-python/tests/test_response_encode.py
@@ -1,6 +1,12 @@
 import unittest
 
-from imap_codec import Encoded, ResponseCodec
+from imap_codec import (
+    Encoded,
+    LineFragment,
+    LiteralFragment,
+    LiteralMode,
+    ResponseCodec,
+)
 
 
 class TestResponseEncode(unittest.TestCase):
@@ -9,7 +15,7 @@ def test_simple_response(self):
         encoded = ResponseCodec.encode(response)
         self.assertIsInstance(encoded, Encoded)
         fragments = list(encoded)
-        self.assertEqual(fragments, [{"Line": {"data": list(b"* SEARCH 1\r\n")}}])
+        self.assertEqual(fragments, [LineFragment(b"* SEARCH 1\r\n")])
 
     def test_simple_response_dump(self):
         response = {"Data": {"Search": [1]}}
@@ -46,9 +52,9 @@ def test_multi_fragment_response(self):
         self.assertEqual(
             fragments,
             [
-                {"Line": {"data": list(b"* 12345 FETCH (BODY[] {5+}\r\n")}},
-                {"Literal": {"data": list(b"ABCDE"), "mode": "NonSync"}},
-                {"Line": {"data": list(b")\r\n")}},
+                LineFragment(b"* 12345 FETCH (BODY[] {5+}\r\n"),
+                LiteralFragment(b"ABCDE", LiteralMode.NonSync),
+                LineFragment(b")\r\n"),
             ],
         )
 
@@ -63,9 +69,7 @@ def test_multi_fragment_response_dump(self):
     def test_multi_fragment_response_dump_remaining(self):
         encoded = ResponseCodec.encode(self._MULTI_FRAGMENT_RESPONSE)
         self.assertIsInstance(encoded, Encoded)
-        self.assertEqual(
-            next(encoded), {"Line": {"data": list(b"* 12345 FETCH (BODY[] {5+}\r\n")}}
-        )
+        self.assertEqual(next(encoded), LineFragment(b"* 12345 FETCH (BODY[] {5+}\r\n"))
         self.assertEqual(
             encoded.dump(),
             b"ABCDE)\r\n",