def test_all_fields_in_document_proto():
""" This tests: all fields are picked from the proto definition """
document_proto_properties = PROTO_TO_PYDANTIC_MODELS.DocumentProto(
).schema(by_alias=False)['definitions']['DocumentProto']['properties']
for i in [
'id',
'content_hash',
'granularity',
'adjacency',
'level_name',
'parent_id',
'chunks',
'weight',
'matches',
'mime_type',
'uri',
'tags',
'location',
'offset',
'embedding',
'score',
'modality',
'evaluations',
]:
assert i in document_proto_properties
document_proto_properties_alias = PROTO_TO_PYDANTIC_MODELS.DocumentProto(
).schema()['definitions']['DocumentProto']['properties']
for i in ['contentHash', 'levelName', 'parentId', 'mimeType']:
assert i in document_proto_properties_alias
def test_oneof_blob():
""" This tests: oneof field is correctly represented as `anyOf` """
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(
blob=PROTO_TO_PYDANTIC_MODELS.NdArrayProto())
assert doc.blob == PROTO_TO_PYDANTIC_MODELS.NdArrayProto()
assert 'text' not in doc.dict()
assert 'buffer' not in doc.dict()
def test_repeated():
""" This tests: repeated fields are represented as `array` """
assert (PROTO_TO_PYDANTIC_MODELS.DenseNdArrayProto().schema()['properties']
['shape']['type'] == 'array')
assert (PROTO_TO_PYDANTIC_MODELS.NamedScoreProto().schema()['definitions']
['NamedScoreProto']['properties']['operands']['type'] == 'array')
assert (PROTO_TO_PYDANTIC_MODELS.DocumentProto().schema()['definitions']
['DocumentProto']['properties']['chunks']['type'] == 'array')
def test_tags_document():
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(hello='world')
assert doc.tags == {'hello': 'world'}
assert MessageToDict(Document(doc.dict()).tags) == {'hello': 'world'}
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(hello='world', tags={'key': 'value'})
assert doc.tags == {'hello': 'world', 'key': 'value'}
assert MessageToDict(Document(doc.dict()).tags) == {
'hello': 'world',
'key': 'value',
}
def test_oneof_validation_error():
""" This tests validation error for invalid fields """
with pytest.raises(pydantic.error_wrappers.ValidationError) as error:
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(text='abc', buffer=b'abc')
assert "only one field among ['buffer', 'blob', 'text']" in str(error.value)
with pytest.raises(pydantic.error_wrappers.ValidationError) as error:
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(
text='abc', buffer=b'abc', blob=PROTO_TO_PYDANTIC_MODELS.NdArrayProto()
)
assert "only one field among ['buffer', 'blob', 'text']" in str(error.value)
def test_enum_definitions():
""" This tests: all enums are defined properly as different levels """
quantization_enum_definition = PROTO_TO_PYDANTIC_MODELS.DocumentProto(
).schema()['definitions']['QuantizationMode']
assert quantization_enum_definition['enum'] == [0, 1, 2]
status_code_enum_definition = PROTO_TO_PYDANTIC_MODELS.StatusProto(
).schema()['definitions']['StatusCode']
assert status_code_enum_definition['enum'] == [0, 1, 2, 3, 4, 5, 6]
command_enum_definition = PROTO_TO_PYDANTIC_MODELS.RequestProto().schema(
)['definitions']['Command']
assert command_enum_definition['enum'] == [0, 1, 3, 4, 5, 6]
def test_timestamp():
""" This tests: google.protobuf.Timestamp are represented as date-time """
assert (
PROTO_TO_PYDANTIC_MODELS.RouteProto().schema(by_alias=False)['properties'][
'start_time'
]['type']
== 'string'
)
assert (
PROTO_TO_PYDANTIC_MODELS.RouteProto().schema(by_alias=False)['properties'][
'start_time'
]['format']
== 'date-time'
)
def test_oneof_buffer():
""" This tests: oneof field is correctly represented as `anyOf` """
doc = PROTO_TO_PYDANTIC_MODELS.DocumentProto(buffer=b'abc')
assert doc.buffer == b'abc'
assert 'text' not in doc.dict()
assert 'blob' not in doc.dict()
def test_struct():
""" This tests: google.protobuf.Struct are represented as `object` """
assert (
PROTO_TO_PYDANTIC_MODELS.DocumentProto().schema()['definitions'][
'DocumentProto'
]['properties']['tags']['type']
== 'object'
)
def test_existing_definitions():
""" This tests: all internal schema definitions are part of parent """
for i in [
'QuantizationMode',
'DenseNdArrayProto',
'SparseNdArrayProto',
'NdArrayProto',
'NamedScoreProto',
'DocumentProto',
]:
assert (i in PROTO_TO_PYDANTIC_MODELS.DocumentProto().schema()
['definitions'].keys())
def test_recursive_schema():
""" This tests: recursive schmea definions are represented properly """
assert PROTO_TO_PYDANTIC_MODELS.NamedScoreProto().schema(
)['definitions']['NamedScoreProto']['properties']['operands']['items'] == {
'$ref': '#/definitions/NamedScoreProto'
}
def test_schema_invocation():
for k, v in PROTO_TO_PYDANTIC_MODELS.items():
v.schema()
v.schema_json()
