fix schema tests

Author: BewareMyPower

pulsar/asyncio.py

@@ -72,7 +72,7 @@ class Producer:
     The Pulsar message producer, used to publish messages on a topic.
     """
 
-    def __init__(self, producer: _pulsar.Producer) -> None:
+    def __init__(self, producer: _pulsar.Producer, schema: pulsar.schema.Schema) -> None:
         """
         Create the producer.
         Users should not call this constructor directly. Instead, create the
@@ -82,8 +82,11 @@ def __init__(self, producer: _pulsar.Producer) -> None:
         ----------
         producer: _pulsar.Producer
             The underlying Producer object from the C extension.
+        schema: pulsar.schema.Schema
+            The schema of the data that will be sent by this producer.
         """
-        self._producer: _pulsar.Producer = producer
+        self._producer = producer
+        self._schema = schema
 
     async def send(self, content: Any) -> pulsar.MessageId:
         """
@@ -105,7 +108,7 @@ async def send(self, content: Any) -> pulsar.MessageId:
         PulsarException
         """
         builder = _pulsar.MessageBuilder()
-        builder.content(content)
+        builder.content(self._schema.encode(content))
         future = asyncio.get_running_loop().create_future()
         self._producer.send_async(builder.build(), functools.partial(_set_future, future))
         msg_id = await future
@@ -454,7 +457,8 @@ def underlying_router(msg, num_partitions):
         self._client.create_producer_async(
             topic, conf, functools.partial(_set_future, future)
         )
-        return Producer(await future)
+        schema.attach_client(self._client)
+        return Producer(await future, schema)
 
     # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-positional-arguments
     async def subscribe(self, topic: Union[str, List[str]],
@@ -632,11 +636,9 @@ async def subscribe(self, topic: Union[str, List[str]],
                 functools.partial(_set_future, future)
             )
         else:
-            raise ValueError(
-                "Argument 'topic' is expected to be of a type between "
-                "(str, list)"
-            )
+            raise ValueError( "Argument 'topic' is expected to be of type 'str' or 'list'")
 
+        schema.attach_client(self._client)
         return Consumer(await future, schema)
 
     async def close(self) -> None:

tests/asyncio_test.py

@@ -39,6 +39,12 @@
     Producer,
     PulsarException,
 )
+from pulsar.schema import (  # pylint: disable=import-error
+    AvroSchema,
+    Integer,
+    Record,
+    String,
+)
 
 SERVICE_URL = 'pulsar://localhost:6650'
 
@@ -253,16 +259,23 @@ async def test_seek_timestamp(self):
         self.assertEqual(msg.data(), b'msg-3')
 
     async def test_schema(self):
+        class ExampleRecord(Record):  # pylint: disable=too-few-public-methods
+            """Example record schema for testing."""
+            str_field = String()
+            int_field = Integer()
+
         topic = f'asyncio-test-schema-{time.time()}'
         producer = await self._client.create_producer(
-                topic, schema=pulsar.schema.StringSchema()
+                topic, schema=AvroSchema(ExampleRecord)
         )
         consumer = await self._client.subscribe(
-            topic, 'sub', schema=pulsar.schema.StringSchema()
+            topic, 'sub', schema=AvroSchema(ExampleRecord)
         )
-        await producer.send('test-message')
+        await producer.send(ExampleRecord(str_field='test', int_field=42))
         msg = await consumer.receive()
-        self.assertEqual(msg.value(), 'test-message')
+        self.assertIsInstance(msg.value(), ExampleRecord)
+        self.assertEqual(msg.value().str_field, 'test')
+        self.assertEqual(msg.value().int_field, 42)
 
 
 if __name__ == '__main__':