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_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_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_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_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_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]