redis-om-python/tests/test_json_model.py
2021-10-13 12:58:10 -07:00

487 lines
No EOL
14 KiB
Python

import abc
import decimal
import datetime
from typing import Optional, List
from unittest import mock
import pytest
import redis
from pydantic import ValidationError
from redis_developer.orm import (
JsonModel,
Field,
)
from redis_developer.orm.model import RedisModelError, QueryNotSupportedError, NotFoundError, embedded
r = redis.Redis()
today = datetime.date.today()
class BaseJsonModel(JsonModel, abc.ABC):
class Meta:
global_key_prefix = "redis-developer"
class EmbeddedJsonModel(BaseJsonModel, abc.ABC):
class Meta:
embedded = True
class Note(EmbeddedJsonModel):
description: str = Field(index=True, full_text_search=True)
created_on: datetime.datetime
class Address(EmbeddedJsonModel):
address_line_1: str
address_line_2: Optional[str]
city: str = Field(index=True)
state: str
country: str
postal_code: str = Field(index=True)
note: Optional[Note]
class Item(EmbeddedJsonModel):
price: decimal.Decimal
name: str = Field(index=True, full_text_search=True)
class Order(EmbeddedJsonModel):
items: List[Item]
total: decimal.Decimal
created_on: datetime.datetime
class Member(BaseJsonModel):
first_name: str = Field(index=True)
last_name: str = Field(index=True)
email: str = Field(index=True)
join_date: datetime.date
age: int = Field(index=True)
# Creates an embedded model.
address: Address
# Creates an embedded list of models.
orders: Optional[List[Order]]
@pytest.fixture()
def address():
yield Address(
address_line_1="1 Main St.",
city="Portland",
state="OR",
country="USA",
postal_code=11111
)
@pytest.fixture()
def members(address):
member1 = Member(
first_name="Andrew",
last_name="Brookins",
email="a@example.com",
age=38,
join_date=today,
address=address
)
member2 = Member(
first_name="Kim",
last_name="Brookins",
email="k@example.com",
age=34,
join_date=today,
address=address
)
member3 = Member(
first_name="Andrew",
last_name="Smith",
email="as@example.com",
age=100,
join_date=today,
address=address
)
member1.save()
member2.save()
member3.save()
yield member1, member2, member3
def test_validates_required_fields(address):
# Raises ValidationError address is required
with pytest.raises(ValidationError):
Member(
first_name="Andrew",
last_name="Brookins",
zipcode="97086",
join_date=today,
)
def test_validates_field(address):
# Raises ValidationError: join_date is not a date
with pytest.raises(ValidationError):
Member(
first_name="Andrew",
last_name="Brookins",
join_date="yesterday",
address=address
)
# Passes validation
def test_validation_passes(address):
member = Member(
first_name="Andrew",
last_name="Brookins",
email="a@example.com",
join_date=today,
age=38,
address=address
)
assert member.first_name == "Andrew"
def test_saves_model_and_creates_pk(address):
member = Member(
first_name="Andrew",
last_name="Brookins",
email="a@example.com",
join_date=today,
age=38,
address=address
)
# Save a model instance to Redis
member.save()
member2 = Member.get(member.pk)
assert member2 == member
assert member2.address == address
@pytest.mark.skip("Not implemented yet")
def test_saves_many(address):
members = [
Member(
first_name="Andrew",
last_name="Brookins",
email="a@example.com",
join_date=today,
address=address,
age=38
),
Member(
first_name="Kim",
last_name="Brookins",
email="k@example.com",
join_date=today,
address=address,
age=34
)
]
Member.add(members)
@pytest.mark.skip("Not ready yet")
def test_updates_a_model(members):
member1, member2, member3 = members
# Or, with an implicit save:
member1.update(last_name="Smith")
assert Member.find(Member.pk == member1.pk).first() == member1
# Or, affecting multiple model instances with an implicit save:
Member.find(Member.last_name == "Brookins").update(last_name="Smith")
results = Member.find(Member.last_name == "Smith")
assert results == members
# Or, updating a field in an embedded model:
member2.update(address__city="Happy Valley")
assert Member.find(Member.pk == member2.pk).first().address.city == "Happy Valley"
def test_paginate_query(members):
member1, member2, member3 = members
actual = Member.find().all(batch_size=1)
assert actual == [member1, member2, member3]
def test_access_result_by_index_cached(members):
member1, member2, member3 = members
query = Member.find().sort_by('age')
# Load the cache, throw away the result.
assert query._model_cache == []
query.execute()
assert query._model_cache == [member2, member1, member3]
# Access an item that should be in the cache.
with mock.patch.object(query.model, 'db') as mock_db:
assert query[0] == member2
assert not mock_db.called
def test_access_result_by_index_not_cached(members):
member1, member2, member3 = members
query = Member.find().sort_by('age')
# Assert that we don't have any models in the cache yet -- we
# haven't made any requests of Redis.
assert query._model_cache == []
assert query[0] == member2
assert query[1] == member1
assert query[2] == member3
def test_exact_match_queries(members):
member1, member2, member3 = members
actual = Member.find(Member.last_name == "Brookins").all()
assert actual == [member1, member2]
actual = Member.find(
(Member.last_name == "Brookins") & ~(Member.first_name == "Andrew")).all()
assert actual == [member2]
actual = Member.find(~(Member.last_name == "Brookins")).all()
assert actual == [member3]
actual = Member.find(Member.last_name != "Brookins").all()
assert actual == [member3]
actual = Member.find(
(Member.last_name == "Brookins") & (Member.first_name == "Andrew")
| (Member.first_name == "Kim")
).all()
assert actual == [member1, member2]
actual = Member.find(Member.first_name == "Kim", Member.last_name == "Brookins").all()
assert actual == [member2]
actual = Member.find(Member.address.city == "Portland").all()
assert actual == [member1, member2, member3]
def test_recursive_query_expression_resolution(members):
member1, member2, member3 = members
actual = Member.find((Member.last_name == "Brookins") | (
Member.age == 100
) & (Member.last_name == "Smith")).all()
assert actual == [member1, member2, member3]
def test_recursive_query_field_resolution(members):
member1, _, _ = members
member1.address.note = Note(description="Weird house",
created_on=datetime.datetime.now())
member1.save()
actual = Member.find(Member.address.note.description == "Weird house").all()
assert actual == [member1]
member1.orders = [
Order(items=[Item(price=10.99, name="Ball")],
total=10.99,
created_on=datetime.datetime.now())
]
member1.save()
actual = Member.find(Member.orders.items.name == "Ball").all()
assert actual == [member1]
assert actual[0].orders[0].items[0].name == "Ball"
def test_full_text_search(members):
member1, member2, _ = members
member1.address.note = Note(description="white house",
created_on=datetime.datetime.now())
member2.address.note = Note(description="blue house",
created_on=datetime.datetime.now())
member1.save()
member2.save()
actual = Member.find(Member.address.note.description % "white").all()
assert actual == [member1]
member1.orders = [
Order(items=[Item(price=10.99, name="balls")],
total=10.99,
created_on=datetime.datetime.now())
]
member2.orders = [
Order(items=[Item(price=10.99, name="white ball")],
total=10.99,
created_on=datetime.datetime.now())
]
member1.save()
member2.save()
actual = Member.find(Member.orders.items.name % "ball").all()
assert actual == [member1, member2]
def test_tag_queries_boolean_logic(members):
member1, member2, member3 = members
actual = Member.find(
(Member.first_name == "Andrew") &
(Member.last_name == "Brookins") | (Member.last_name == "Smith")).all()
assert actual == [member1, member3]
def test_tag_queries_punctuation(address):
member1 = Member(
first_name="Andrew, the Michael",
last_name="St. Brookins-on-Pier",
email="a|b@example.com", # NOTE: This string uses the TAG field separator.
age=38,
join_date=today,
address=address
)
member1.save()
member2 = Member(
first_name="Bob",
last_name="the Villain",
email="a|villain@example.com", # NOTE: This string uses the TAG field separator.
age=38,
join_date=today,
address=address
)
member2.save()
assert Member.find(Member.first_name == "Andrew, the Michael").first() == member1
assert Member.find(Member.last_name == "St. Brookins-on-Pier").first() == member1
# Notice that when we index and query multiple values that use the internal
# TAG separator for single-value exact-match fields, like an indexed string,
# the queries will succeed. We apply a workaround that queries for the union
# of the two values separated by the tag separator.
assert Member.find(Member.email == "a|b@example.com").all() == [member1]
assert Member.find(Member.email == "a|villain@example.com").all() == [member2]
def test_tag_queries_negation(members):
member1, member2, member3 = members
"""
┌first_name
NOT EQ┤
└Andrew
"""
query = Member.find(
~(Member.first_name == "Andrew")
)
assert query.all() == [member2]
"""
┌first_name
┌NOT EQ┤
| └Andrew
AND┤
| ┌last_name
└EQ┤
└Brookins
"""
query = Member.find(
~(Member.first_name == "Andrew") & (Member.last_name == "Brookins")
)
assert query.all() == [member2]
"""
┌first_name
┌NOT EQ┤
| └Andrew
AND┤
| ┌last_name
| ┌EQ┤
| | └Brookins
└OR┤
| ┌last_name
└EQ┤
└Smith
"""
query = Member.find(
~(Member.first_name == "Andrew") &
((Member.last_name == "Brookins") | (Member.last_name == "Smith")))
assert query.all() == [member2]
"""
┌first_name
┌NOT EQ┤
| └Andrew
┌AND┤
| | ┌last_name
| └EQ┤
| └Brookins
OR┤
| ┌last_name
└EQ┤
└Smith
"""
query = Member.find(
~(Member.first_name == "Andrew") &
(Member.last_name == "Brookins") | (Member.last_name == "Smith"))
assert query.all() == [member2, member3]
actual = Member.find(
(Member.first_name == "Andrew") & ~(Member.last_name == "Brookins")).all()
assert actual == [member3]
def test_numeric_queries(members):
member1, member2, member3 = members
actual = Member.find(Member.age == 34).all()
assert actual == [member2]
actual = Member.find(Member.age > 34).all()
assert actual == [member1, member3]
actual = Member.find(Member.age < 35).all()
assert actual == [member2]
actual = Member.find(Member.age <= 34).all()
assert actual == [member2]
actual = Member.find(Member.age >= 100).all()
assert actual == [member3]
actual = Member.find(~(Member.age == 100)).all()
assert actual == [member1, member2]
def test_sorting(members):
member1, member2, member3 = members
actual = Member.find(Member.age > 34).sort_by('age').all()
assert actual == [member1, member3]
actual = Member.find(Member.age > 34).sort_by('-age').all()
assert actual == [member3, member1]
with pytest.raises(QueryNotSupportedError):
# This field does not exist.
Member.find().sort_by('not-a-real-field').all()
with pytest.raises(QueryNotSupportedError):
# This field is not sortable.
Member.find().sort_by('join_date').all()
def test_not_found():
with pytest.raises(NotFoundError):
# This ID does not exist.
Member.get(1000)
def test_schema():
assert Member.redisearch_schema() == "ON JSON PREFIX 1 redis-developer:tests.test_json_model.Member: SCHEMA $.pk AS pk TAG SEPARATOR | SORTABLE $.first_name AS first_name TAG SEPARATOR | SORTABLE $.last_name AS last_name TAG SEPARATOR | SORTABLE $.email AS email TAG SEPARATOR | SORTABLE $.age AS age NUMERIC SORTABLE $.address.pk AS address_pk TAG SEPARATOR | SORTABLE $.address.city AS address_city TAG SEPARATOR | SORTABLE $.address.postal_code AS address_postal_code TAG SEPARATOR | SORTABLE $.address.note.pk AS address_note_pk TAG SEPARATOR | SORTABLE $.address.note.description AS address_note_description TAG SEPARATOR | $.address.note.description AS address_note_description_fts TEXT SORTABLE $.orders[*].pk AS orders_pk TAG SEPARATOR | SORTABLE $.orders[*].items[*].pk AS orders_items_pk TAG SEPARATOR | SORTABLE $.orders[*].items[*].name AS orders_items_name TAG SEPARATOR | $.orders[*].items[*].name AS orders_items_name_fts TEXT SORTABLE"