Add tests for jsonschema

Ron Serruya · Ron Serruya · commit 4e3f26778463 · 2021-11-18T15:09:48.000+02:00
diff --git a/tests/integration/kafka/test_kafka_integration.py b/tests/integration/kafka/test_kafka_integration.py
@@ -7,10 +7,12 @@
 
 from aws_schema_registry import DataAndSchema, SchemaRegistryClient
 from aws_schema_registry.avro import AvroSchema
+from aws_schema_registry.jsonschema import JsonSchema
 from aws_schema_registry.adapter.kafka import (
     KafkaDeserializer, KafkaSerializer
 )
-from aws_schema_registry.naming import record_name_strategy
+from aws_schema_registry.naming import (record_name_strategy,
+                                        topic_name_strategy)
 
 BOOTSTRAP_STRING = '127.0.0.1:9092'
 
@@ -24,6 +26,8 @@
     SCHEMA_V1 = AvroSchema(f.read())
 with open(os.path.join(os.path.dirname(__file__), 'user.v2.avsc'), 'r') as f:
     SCHEMA_V2 = AvroSchema(f.read())
+with open(os.path.join(os.path.dirname(__file__), 'user.json'), 'r') as f:
+    SCHEMA_JSON = JsonSchema(f.read())
 
 PRODUCER_PROPERTIES = {
     'bootstrap_servers': BOOTSTRAP_STRING,
@@ -65,12 +69,25 @@ def test_produce_consume_with_ser_de_schema_registry(
     serializer = KafkaSerializer(
         client, schema_naming_strategy=record_name_strategy
     )
+
+    # jsonschema has no fqn, so we use topic_name_strategy for it
+    # (which also requires a separate producer)
+    json_serializer = KafkaSerializer(
+        client, schema_naming_strategy=topic_name_strategy
+    )
+
     deserializer = KafkaDeserializer(client)
 
     producer = KafkaProducer(
         value_serializer=serializer,
         **PRODUCER_PROPERTIES
     )
+
+    json_producer = KafkaProducer(
+        value_serializer=json_serializer,
+        **PRODUCER_PROPERTIES
+    )
+
     data1 = {
         'name': 'John Doe',
         'favorite_number': 6,
@@ -85,6 +102,13 @@ def test_produce_consume_with_ser_de_schema_registry(
     }
     producer.send(topic, DataAndSchema(data2, SCHEMA_V2))
 
+    data3 = {
+        'name': 'John Doe',
+        'favorite_number': 6,
+        'favorite_colors': ['red', 'blue', "yello"]
+    }
+    json_producer.send(topic, DataAndSchema(data3, SCHEMA_JSON))
+
     consumer = KafkaConsumer(
         topic,
         value_deserializer=deserializer,
@@ -93,8 +117,10 @@ def test_produce_consume_with_ser_de_schema_registry(
     batch = consumer.poll(timeout_ms=1000)
     assert len(batch) == 1
     messages = batch[list(batch.keys())[0]]
-    assert len(messages) == 2
+    assert len(messages) == 3
     assert messages[0].value.data == data1
     assert messages[0].value.schema == SCHEMA_V1
     assert messages[1].value.data == data2
     assert messages[1].value.schema == SCHEMA_V2
+    assert messages[2].value.data == data3
+    assert messages[2].value.schema == SCHEMA_JSON
diff --git a/tests/integration/kafka/user.json b/tests/integration/kafka/user.json
@@ -0,0 +1,23 @@
+{
+  "$schema": "http://json-schema.org/draft-04/schema#",
+  "type": "object",
+  "properties": {
+    "name": {
+      "type": "string"
+    },
+    "favorite_number": {
+      "type": "integer"
+    },
+    "favorite_colors": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      }
+    }
+  },
+  "required": [
+    "name",
+    "favorite_number",
+    "favorite_colors"
+  ]
+}
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -10,9 +10,27 @@
 
 REGISTRY_NAME = 'user-topic'
 SCHEMA_NAME = 'User-Topic'
+JSON_SCHEMA_NAME = 'User-Topic'
 SCHEMA_ARN = f'arn:aws:glue:us-west-2:123:schema/{REGISTRY_NAME}/{SCHEMA_NAME}'
 SCHEMA_VERSION_ID = UUID('b7b4a7f0-9c96-4e4a-a687-fb5de9ef0c63')
+JSON_SCHEMA_VERSION_ID = UUID('98718bb6-ca2a-4ac6-b841-748cab68b1b1')
 SCHEMA_DEF = '{"name": "Test", "type": "record", "fields": []}'
+JSON_SCHEMA_DEF = """{
+      "$schema": "http://json-schema.org/draft-04/schema#",
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "age": {
+          "type": "integer"
+        }
+      },
+      "required": [
+        "name",
+        "age"
+      ]
+    }"""
 
 METADATA = {
     'event-source-1': 'topic1',
@@ -97,6 +115,23 @@ def test_get_or_register_schema_version_creates_schema(client, glue_client):
 
     assert version.version_id == SCHEMA_VERSION_ID
 
+    glue_client.get_schema_version = Mock(return_value={
+        'SchemaVersionId': str(SCHEMA_VERSION_ID),
+        'SchemaDefinition': SCHEMA_DEF,
+        'DataFormat': 'JSON',
+        'SchemaArn': SCHEMA_ARN,
+        'VersionNumber': 123,
+        'Status': 'AVAILABLE'
+    })
+
+    version = client.get_or_register_schema_version(
+        definition=SCHEMA_DEF,
+        schema_name=SCHEMA_NAME,
+        data_format='JSON'
+    )
+
+    assert version.version_id == SCHEMA_VERSION_ID
+
 
 def test_get_or_register_schema_version_registers_version(
     client, glue_client
@@ -128,16 +163,22 @@ def test_get_or_register_schema_version_registers_version(
     assert version.version_id == SCHEMA_VERSION_ID
 
 
-def test_register_schema_version(client, glue_client):
+@pytest.mark.parametrize(
+    "schema_def,schema_name,schema_ver_id",
+    [(SCHEMA_DEF, SCHEMA_NAME, SCHEMA_VERSION_ID),
+     (JSON_SCHEMA_DEF, JSON_SCHEMA_NAME, JSON_SCHEMA_VERSION_ID)])
+def test_register_schema_version(client, glue_client,
+                                 schema_name, schema_def, schema_ver_id):
+    print(schema_name, schema_def, schema_ver_id)
     glue_client.register_schema_version = Mock(return_value={
-        'SchemaVersionId': str(SCHEMA_VERSION_ID),
+        'SchemaVersionId': str(schema_ver_id),
         'VersionNumber': 1,
         'Status': 'AVAILABLE'
     })
 
-    version_id = client.register_schema_version(SCHEMA_DEF, SCHEMA_NAME)
+    version_id = client.register_schema_version(schema_def, schema_name)
 
-    assert version_id == SCHEMA_VERSION_ID
+    assert version_id == schema_ver_id
 
 
 def test_wait_for_schema_evolution_check_to_complete(client, glue_client):
@@ -200,16 +241,17 @@ def _make_put_schema_version_metadata_response(
     }
 
 
-def test_create_schema(client, glue_client):
+@pytest.mark.parametrize("data_format", ["AVRO", "JSON"])
+def test_create_schema(client, glue_client, data_format):
     schema_version_id = uuid4()
     glue_client.create_schema = Mock(return_value={
         'SchemaName': SCHEMA_NAME,
-        'DataFormat': 'AVRO',
+        'DataFormat': data_format,
         'SchemaVersionId': str(schema_version_id)
     })
 
     version_id = client.create_schema(
-        SCHEMA_NAME, 'AVRO', SCHEMA_DEF
+        SCHEMA_NAME, data_format, SCHEMA_DEF
     )
 
     assert version_id == schema_version_id
diff --git a/tests/test_jsonschema.py b/tests/test_jsonschema.py
@@ -0,0 +1,87 @@
+import re
+import pytest
+
+from aws_schema_registry import ValidationError
+from aws_schema_registry.jsonschema import JsonSchema
+
+
+def test_readwrite():
+    s = JsonSchema("""{
+          "$schema": "http://json-schema.org/draft-04/schema#",
+          "type": "object",
+          "properties": {
+            "name": {
+              "type": "string"
+            },
+            "age": {
+              "type": "integer"
+            }
+          },
+          "required": [
+            "name",
+            "age"
+          ]
+        }""")
+
+    d = {
+        'name': 'Yoda',
+        'age': 900
+    }
+
+    assert s.read(s.write(d)) == d
+
+
+def test_validation():
+    s = JsonSchema("""{
+      "$schema": "http://json-schema.org/draft-04/schema#",
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "age": {
+          "type": "integer"
+        }
+      },
+      "required": [
+        "name",
+        "age"
+      ]
+    }""")
+
+    with pytest.raises(ValidationError, match=re.escape(
+        "data must contain ['name', 'age'] properties"
+    )):
+
+        s.validate({'name': 'Obi-Wan'})
+    with pytest.raises(ValidationError, match=re.escape(
+        "data.name must be string"
+    )):
+
+        s.validate({'name': 1, 'age': 2})
+
+    s.validate({'name': 'Jar Jar', 'age': 42, 'sith': True})
+
+    s = JsonSchema("""{
+          "$schema": "http://json-schema.org/draft-04/schema#",
+          "type": "object",
+          "properties": {
+            "name": {
+              "type": "string"
+            },
+            "age": {
+              "type": "integer"
+            }
+          },
+          "required": [
+            "name",
+            "age"
+          ],
+          "additionalProperties": false
+        }""")
+
+    with pytest.raises(ValidationError, match=re.escape(
+        "data must not contain {'sith'} properties"
+    )):
+
+        s.validate({'name': 'Jar Jar', 'age': 42, 'sith': True})
