Add unit tests for async SR client and async serdes (#1994)

* Add unit tests for async SR client and async serdes

* fix flake8 issues

* fix test_delivery_report_serialization

* run unasync

Author: rohitsanj

src/confluent_kafka/schema_registry/_async/avro.py

@@ -558,7 +558,7 @@ async def __deserialize(
                 latest_schema = await self._get_reader_schema(subject)
 
         if latest_schema is not None:
-            migrations = self._get_migrations(subject, writer_schema_raw, latest_schema, None)
+            migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema
             reader_schema = await self._get_parsed_schema(latest_schema.schema)
         elif self._schema is not None:

src/confluent_kafka/schema_registry/_async/json_schema.py

@@ -598,7 +598,7 @@ async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = N
             writer_schema, writer_ref_registry = None, None
 
         if latest_schema is not None:
-            migrations = self._get_migrations(subject, writer_schema_raw, latest_schema, None)
+            migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema
             reader_schema, reader_ref_registry = await self._get_parsed_schema(latest_schema.schema)
         elif self._schema is not None:

src/confluent_kafka/schema_registry/_async/mock_schema_registry_client.py

@@ -0,0 +1,264 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2024 Confluent Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import uuid
+from collections import defaultdict
+from threading import Lock
+from typing import List, Dict, Optional
+
+from .schema_registry_client import AsyncSchemaRegistryClient
+from ..common.schema_registry_client import RegisteredSchema, Schema, ServerConfig
+from ..error import SchemaRegistryError
+
+
+class _SchemaStore(object):
+
+    def __init__(self):
+        self.lock = Lock()
+        self.max_id = 0
+        self.schema_id_index = {}
+        self.schema_guid_index = {}
+        self.schema_index = {}
+        self.subject_schemas = defaultdict(set)
+
+    def set(self, registered_schema: RegisteredSchema) -> RegisteredSchema:
+        with self.lock:
+            self.max_id += 1
+            rs = RegisteredSchema(
+                schema_id=self.max_id,
+                guid=registered_schema.guid,
+                schema=registered_schema.schema,
+                subject=registered_schema.subject,
+                version=registered_schema.version
+            )
+            self.schema_id_index[rs.schema_id] = rs
+            self.schema_guid_index[rs.guid] = rs
+            self.schema_index[rs.schema] = rs.schema_id
+            self.subject_schemas[rs.subject].add(rs)
+            return rs
+
+    def get_schema(self, schema_id: int) -> Optional[Schema]:
+        with self.lock:
+            rs = self.schema_id_index.get(schema_id, None)
+            return rs.schema if rs else None
+
+    def get_schema_by_guid(self, guid: str) -> Optional[Schema]:
+        with self.lock:
+            rs = self.schema_guid_index.get(guid, None)
+            return rs.schema if rs else None
+
+    def get_registered_schema_by_schema(
+        self,
+        subject_name: str,
+        schema: Schema
+    ) -> Optional[RegisteredSchema]:
+        with self.lock:
+            if subject_name in self.subject_schemas:
+                for rs in self.subject_schemas[subject_name]:
+                    if rs.schema == schema:
+                        return rs
+            return None
+
+    def get_version(self, subject_name: str, version: int) -> Optional[RegisteredSchema]:
+        with self.lock:
+            if subject_name in self.subject_schemas:
+                for rs in self.subject_schemas[subject_name]:
+                    if rs.version == version:
+                        return rs
+            return None
+
+    def get_latest_version(self, subject_name: str) -> Optional[RegisteredSchema]:
+        with self.lock:
+            if subject_name in self.subject_schemas:
+                latest_version = 0
+                latest_schema = None
+                for rs in self.subject_schemas[subject_name]:
+                    if rs.version > latest_version:
+                        latest_version = rs.version
+                        latest_schema = rs
+                return latest_schema
+            return None
+
+    def get_latest_with_metadata(
+        self, subject_name: str,
+        metadata: Dict[str, str]
+    ) -> Optional[RegisteredSchema]:
+        with self.lock:
+            if subject_name in self.subject_schemas:
+                rs: RegisteredSchema
+                for rs in self.subject_schemas[subject_name]:
+                    if (rs.schema
+                            and rs.schema.metadata
+                            and rs.schema.metadata.properties
+                            and metadata.items() <= rs.schema.metadata.properties.properties.items()):
+                        return rs
+            return None
+
+    def get_subjects(self) -> List[str]:
+        with self.lock:
+            return list(self.subject_schemas.keys())
+
+    def get_versions(self, subject_name: str) -> List[int]:
+        with self.lock:
+            if subject_name in self.subject_schemas:
+                return [rs.version for rs in self.subject_schemas[subject_name]]
+            return []
+
+    def remove_by_schema(self, registered_schema: RegisteredSchema):
+        with self.lock:
+            subject_name = registered_schema.subject
+            if subject_name in self.subject_schemas:
+                self.subject_schemas[subject_name].remove(registered_schema)
+
+    def remove_by_subject(self, subject_name: str) -> List[int]:
+        with self.lock:
+            versions = []
+            if subject_name in self.subject_schemas:
+                for rs in self.subject_schemas[subject_name]:
+                    versions.append(rs.version)
+                    schema_id = self.schema_index.pop(rs.schema, None)
+                    if schema_id is not None:
+                        self.schema_id_index.pop(schema_id, None)
+
+                del self.subject_schemas[subject_name]
+            return versions
+
+    def clear(self):
+        with self.lock:
+            self.schema_id_index.clear()
+            self.schema_guid_index.clear()
+            self.schema_index.clear()
+            self.subject_schemas.clear()
+
+
+class AsyncMockSchemaRegistryClient(AsyncSchemaRegistryClient):
+
+    def __init__(self, conf: dict):
+        super().__init__(conf)
+        self._store = _SchemaStore()
+
+    async def register_schema(
+        self, subject_name: str, schema: 'Schema',
+        normalize_schemas: bool = False
+    ) -> int:
+        registered_schema = await self.register_schema_full_response(subject_name, schema, normalize_schemas)
+        return registered_schema.schema_id
+
+    async def register_schema_full_response(
+        self, subject_name: str, schema: 'Schema',
+        normalize_schemas: bool = False
+    ) -> 'RegisteredSchema':
+        registered_schema = self._store.get_registered_schema_by_schema(subject_name, schema)
+        if registered_schema is not None:
+            return registered_schema
+
+        latest_schema = self._store.get_latest_version(subject_name)
+        latest_version = 1 if latest_schema is None else latest_schema.version + 1
+
+        registered_schema = RegisteredSchema(
+            schema_id=1,
+            guid=str(uuid.uuid4()),
+            schema=schema,
+            subject=subject_name,
+            version=latest_version
+        )
+
+        registered_schema = self._store.set(registered_schema)
+
+        return registered_schema
+
+    async def get_schema(
+        self, schema_id: int, subject_name: Optional[str] = None,
+        fmt: Optional[str] = None
+    ) -> 'Schema':
+        schema = self._store.get_schema(schema_id)
+        if schema is not None:
+            return schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def get_schema_by_guid(
+        self, guid: str, fmt: Optional[str] = None
+    ) -> 'Schema':
+        schema = self._store.get_schema_by_guid(guid)
+        if schema is not None:
+            return schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def lookup_schema(
+        self, subject_name: str, schema: 'Schema',
+        normalize_schemas: bool = False, deleted: bool = False
+    ) -> 'RegisteredSchema':
+
+        registered_schema = self._store.get_registered_schema_by_schema(subject_name, schema)
+        if registered_schema is not None:
+            return registered_schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def get_subjects(self) -> List[str]:
+        return self._store.get_subjects()
+
+    async def delete_subject(self, subject_name: str, permanent: bool = False) -> List[int]:
+        return self._store.remove_by_subject(subject_name)
+
+    async def get_latest_version(self, subject_name: str, fmt: Optional[str] = None) -> 'RegisteredSchema':
+        registered_schema = self._store.get_latest_version(subject_name)
+        if registered_schema is not None:
+            return registered_schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def get_latest_with_metadata(
+        self, subject_name: str, metadata: Dict[str, str],
+        deleted: bool = False, fmt: Optional[str] = None
+    ) -> 'RegisteredSchema':
+        registered_schema = self._store.get_latest_with_metadata(subject_name, metadata)
+        if registered_schema is not None:
+            return registered_schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def get_version(
+        self, subject_name: str, version: int,
+        deleted: bool = False, fmt: Optional[str] = None
+    ) -> 'RegisteredSchema':
+        registered_schema = self._store.get_version(subject_name, version)
+        if registered_schema is not None:
+            return registered_schema
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def get_versions(self, subject_name: str) -> List[int]:
+        return self._store.get_versions(subject_name)
+
+    async def delete_version(self, subject_name: str, version: int, permanent: bool = False) -> int:
+        registered_schema = self._store.get_version(subject_name, version)
+        if registered_schema is not None:
+            self._store.remove_by_schema(registered_schema)
+            return registered_schema.schema_id
+
+        raise SchemaRegistryError(404, 40400, "Schema Not Found")
+
+    async def set_config(
+        self, subject_name: Optional[str] = None, config: 'ServerConfig' = None  # noqa F821
+    ) -> 'ServerConfig':  # noqa F821
+        return None
+
+    async def get_config(self, subject_name: Optional[str] = None) -> 'ServerConfig':  # noqa F821
+        return None

src/confluent_kafka/schema_registry/_async/protobuf.py

@@ -597,7 +597,7 @@ async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = N
             writer_schema = None
 
         if latest_schema is not None:
-            migrations = self._get_migrations(subject, writer_schema_raw, latest_schema, None)
+            migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema
             fd_proto, pool = await self._get_parsed_schema(latest_schema.schema)
             reader_schema = pool.FindFileByName(fd_proto.name)

src/confluent_kafka/schema_registry/_async/schema_registry_client.py

@@ -1158,8 +1158,8 @@ def clear_caches(self):
 
     @staticmethod
     def new_client(conf: dict) -> 'AsyncSchemaRegistryClient':
-        from confluent_kafka.schema_registry.mock_schema_registry_client import MockSchemaRegistryClient
+        from .mock_schema_registry_client import AsyncMockSchemaRegistryClient
         url = conf.get("url")
         if url.startswith("mock://"):
-            return MockSchemaRegistryClient(conf)
+            return AsyncMockSchemaRegistryClient(conf)
         return AsyncSchemaRegistryClient(conf)

src/confluent_kafka/schema_registry/_async/serde.py

@@ -216,7 +216,7 @@ async def _get_migrations(
         self, subject: str, source_info: Schema,
         target: RegisteredSchema, fmt: Optional[str]
     ) -> List[Migration]:
-        source = self._registry.lookup_schema(subject, source_info, False, True)
+        source = await self._registry.lookup_schema(subject, source_info, False, True)
         migrations = []
         if source.version < target.version:
             migration_mode = RuleMode.UPGRADE

src/confluent_kafka/schema_registry/mock_schema_registry_client.py renamed to src/confluent_kafka/schema_registry/_sync/mock_schema_registry_client.py

@@ -20,8 +20,9 @@
 from threading import Lock
 from typing import List, Dict, Optional
 
-from . import SchemaRegistryClient, RegisteredSchema, Schema
-from .error import SchemaRegistryError
+from .schema_registry_client import SchemaRegistryClient
+from ..common.schema_registry_client import RegisteredSchema, Schema, ServerConfig
+from ..error import SchemaRegistryError
 
 
 class _SchemaStore(object):

src/confluent_kafka/schema_registry/_sync/schema_registry_client.py

@@ -1158,7 +1158,7 @@ def clear_caches(self):
 
     @staticmethod
     def new_client(conf: dict) -> 'SchemaRegistryClient':
-        from confluent_kafka.schema_registry.mock_schema_registry_client import MockSchemaRegistryClient
+        from .mock_schema_registry_client import MockSchemaRegistryClient
         url = conf.get("url")
         if url.startswith("mock://"):
             return MockSchemaRegistryClient(conf)

tests/integration/schema_registry/_async/test_avro_serializers.py

@@ -278,8 +278,8 @@ async def test_delivery_report_serialization(kafka_cluster, load_file, avsc, dat
     producer = kafka_cluster.async_producer(value_serializer=value_serializer)
 
     async def assert_cb(err, msg):
-        actual = value_deserializer(msg.value(),
-                                    SerializationContext(topic, MessageField.VALUE, msg.headers()))
+        actual = await value_deserializer(
+            msg.value(), SerializationContext(topic, MessageField.VALUE, msg.headers()))
 
         if record_type == "record":
             assert [v == actual[k] for k, v in data.items()]

tests/integration/schema_registry/_sync/test_avro_serializers.py

@@ -278,8 +278,8 @@ def test_delivery_report_serialization(kafka_cluster, load_file, avsc, data, rec
     producer = kafka_cluster.producer(value_serializer=value_serializer)
 
     def assert_cb(err, msg):
-        actual = value_deserializer(msg.value(),
-                                    SerializationContext(topic, MessageField.VALUE, msg.headers()))
+        actual = value_deserializer(
+            msg.value(), SerializationContext(topic, MessageField.VALUE, msg.headers()))
 
         if record_type == "record":
             assert [v == actual[k] for k, v in data.items()]