Retain header on decode; test header+msg equality

Author: dpkp

kafka/protocol/api.py

@@ -58,6 +58,11 @@ def API_VERSION(self):
     def to_object(self):
         return _to_object(self.SCHEMA, self)
 
+    @classmethod
+    @abc.abstractmethod
+    def is_request(cls):
+        pass
+
     @property
     def header(self):
         return self._header
@@ -88,16 +93,26 @@ def decode(cls, data, header=False, framed=False):
             return super().decode(data)
         if isinstance(data, bytes):
             data = BytesIO(data)
-        ret = []
         if framed:
-            ret.append(Int32.decode(data))
+            size = Int32.decode(data)
         if header:
-            ret.append(cls.parse_header(data))
-        ret.append(super().decode(data))
-        return tuple(ret)
+            hdr = cls.parse_header(data)
+        else:
+            hdr = None
+        ret = super().decode(data)
+        if hdr is not None:
+            ret._header = hdr
+        return ret
+
+    def __eq__(self, other):
+        return self._header == other._header and super().__eq__(other)
 
 
 class Request(RequestResponse):
+    @classmethod
+    def is_request(cls):
+        return True
+
     def expect_response(self):
         """Override this method if an api request does not always generate a response"""
         return True
@@ -115,13 +130,12 @@ def header_class(cls):
         else:
             return RequestHeader
 
-    def encode(self, header=False, framed=False, correlation_id=None, client_id=None, **kwargs):
-        if header and self.header is None:
-            self.with_header(correlation_id=correlation_id, client_id=client_id)
-        return super().encode(header=header, framed=framed)
-
 
 class Response(RequestResponse):
+    @classmethod
+    def is_request(cls):
+        return False
+
     def with_header(self, correlation_id=0):
         if self.FLEXIBLE_VERSION:
             self._header = self.header_class()(correlation_id, {})
@@ -135,11 +149,6 @@ def header_class(cls):
         else:
             return ResponseHeader
 
-    def encode(self, header=False, framed=False, correlation_id=None, **kwargs):
-        if header and self.header is None:
-            self.with_header(correlation_id=correlation_id)
-        return super().encode(header=header, framed=framed)
-
 
 def _to_object(schema, data):
     obj = {}

kafka/protocol/struct.py

@@ -54,6 +54,8 @@ def __hash__(self):
         return hash(self.encode())
 
     def __eq__(self, other):
+        if not isinstance(other, Struct):
+            return False
         if self.SCHEMA != other.SCHEMA:
             return False
         for attr in self.SCHEMA.names:

test/protocol/test_api_versions.py

@@ -65,5 +65,9 @@
 
 @pytest.mark.parametrize('msg, encoded', TEST_CASES)
 def test_parse(msg, encoded):
-    assert msg.encode(correlation_id=1, client_id='_internal_client_kYVL', header=True, framed=True) == encoded
-    assert msg.decode(encoded, header=True, framed=True)[2] == msg
+    if msg.is_request():
+        msg.with_header(correlation_id=1, client_id='_internal_client_kYVL')
+    else:
+        msg.with_header(correlation_id=1)
+    assert msg.encode(header=True, framed=True) == encoded
+    assert msg.decode(encoded, header=True, framed=True) == msg