453 lines
10 KiB
Python
453 lines
10 KiB
Python
from asyncio import gather
|
|
from collections import namedtuple
|
|
from functools import partial
|
|
from unittest.mock import Mock
|
|
|
|
from graphene.utils.dataloader import DataLoader
|
|
from pytest import mark, raises
|
|
|
|
from graphene import ObjectType, String, Schema, Field, List
|
|
|
|
CHARACTERS = {
|
|
"1": {"name": "Luke Skywalker", "sibling": "3"},
|
|
"2": {"name": "Darth Vader", "sibling": None},
|
|
"3": {"name": "Leia Organa", "sibling": "1"},
|
|
}
|
|
|
|
get_character = Mock(side_effect=lambda character_id: CHARACTERS[character_id])
|
|
|
|
|
|
class CharacterType(ObjectType):
|
|
name = String()
|
|
sibling = Field(lambda: CharacterType)
|
|
|
|
async def resolve_sibling(character, info):
|
|
if character["sibling"]:
|
|
return await info.context.character_loader.load(character["sibling"])
|
|
return None
|
|
|
|
|
|
class Query(ObjectType):
|
|
skywalker_family = List(CharacterType)
|
|
|
|
async def resolve_skywalker_family(_, info):
|
|
return await info.context.character_loader.load_many(["1", "2", "3"])
|
|
|
|
|
|
mock_batch_load_fn = Mock(
|
|
side_effect=lambda character_ids: [get_character(id) for id in character_ids]
|
|
)
|
|
|
|
|
|
class CharacterLoader(DataLoader):
|
|
async def batch_load_fn(self, character_ids):
|
|
return mock_batch_load_fn(character_ids)
|
|
|
|
|
|
Context = namedtuple("Context", "character_loader")
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_basic_dataloader():
|
|
schema = Schema(query=Query)
|
|
|
|
character_loader = CharacterLoader()
|
|
context = Context(character_loader=character_loader)
|
|
|
|
query = """
|
|
{
|
|
skywalkerFamily {
|
|
name
|
|
sibling {
|
|
name
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
|
|
result = await schema.execute_async(query, context=context)
|
|
|
|
assert not result.errors
|
|
assert result.data == {
|
|
"skywalkerFamily": [
|
|
{"name": "Luke Skywalker", "sibling": {"name": "Leia Organa"}},
|
|
{"name": "Darth Vader", "sibling": None},
|
|
{"name": "Leia Organa", "sibling": {"name": "Luke Skywalker"}},
|
|
]
|
|
}
|
|
|
|
assert mock_batch_load_fn.call_count == 1
|
|
assert get_character.call_count == 3
|
|
|
|
|
|
def id_loader(**options):
|
|
load_calls = []
|
|
|
|
async def default_resolve(x):
|
|
return x
|
|
|
|
resolve = options.pop("resolve", default_resolve)
|
|
|
|
async def fn(keys):
|
|
load_calls.append(keys)
|
|
return await resolve(keys)
|
|
# return keys
|
|
|
|
identity_loader = DataLoader(fn, **options)
|
|
return identity_loader, load_calls
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_build_a_simple_data_loader():
|
|
async def call_fn(keys):
|
|
return keys
|
|
|
|
identity_loader = DataLoader(call_fn)
|
|
|
|
promise1 = identity_loader.load(1)
|
|
|
|
value1 = await promise1
|
|
assert value1 == 1
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_can_build_a_data_loader_from_a_partial():
|
|
value_map = {1: "one"}
|
|
|
|
async def call_fn(context, keys):
|
|
return [context.get(key) for key in keys]
|
|
|
|
partial_fn = partial(call_fn, value_map)
|
|
identity_loader = DataLoader(partial_fn)
|
|
|
|
promise1 = identity_loader.load(1)
|
|
|
|
value1 = await promise1
|
|
assert value1 == "one"
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_supports_loading_multiple_keys_in_one_call():
|
|
async def call_fn(keys):
|
|
return keys
|
|
|
|
identity_loader = DataLoader(call_fn)
|
|
|
|
promise_all = identity_loader.load_many([1, 2])
|
|
|
|
values = await promise_all
|
|
assert values == [1, 2]
|
|
|
|
promise_all = identity_loader.load_many([])
|
|
|
|
values = await promise_all
|
|
assert values == []
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_batches_multiple_requests():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
promise1 = identity_loader.load(1)
|
|
promise2 = identity_loader.load(2)
|
|
|
|
p = gather(promise1, promise2)
|
|
|
|
value1, value2 = await p
|
|
|
|
assert value1 == 1
|
|
assert value2 == 2
|
|
|
|
assert load_calls == [[1, 2]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_batches_multiple_requests_with_max_batch_sizes():
|
|
identity_loader, load_calls = id_loader(max_batch_size=2)
|
|
|
|
promise1 = identity_loader.load(1)
|
|
promise2 = identity_loader.load(2)
|
|
promise3 = identity_loader.load(3)
|
|
|
|
p = gather(promise1, promise2, promise3)
|
|
|
|
value1, value2, value3 = await p
|
|
|
|
assert value1 == 1
|
|
assert value2 == 2
|
|
assert value3 == 3
|
|
|
|
assert load_calls == [[1, 2], [3]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_coalesces_identical_requests():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
promise1 = identity_loader.load(1)
|
|
promise2 = identity_loader.load(1)
|
|
|
|
assert promise1 == promise2
|
|
p = gather(promise1, promise2)
|
|
|
|
value1, value2 = await p
|
|
|
|
assert value1 == 1
|
|
assert value2 == 1
|
|
|
|
assert load_calls == [[1]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_caches_repeated_requests():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
a, b = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a == "A"
|
|
assert b == "B"
|
|
|
|
assert load_calls == [["A", "B"]]
|
|
|
|
a2, c = await gather(identity_loader.load("A"), identity_loader.load("C"))
|
|
|
|
assert a2 == "A"
|
|
assert c == "C"
|
|
|
|
assert load_calls == [["A", "B"], ["C"]]
|
|
|
|
a3, b2, c2 = await gather(
|
|
identity_loader.load("A"), identity_loader.load("B"), identity_loader.load("C")
|
|
)
|
|
|
|
assert a3 == "A"
|
|
assert b2 == "B"
|
|
assert c2 == "C"
|
|
|
|
assert load_calls == [["A", "B"], ["C"]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_clears_single_value_in_loader():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
a, b = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a == "A"
|
|
assert b == "B"
|
|
|
|
assert load_calls == [["A", "B"]]
|
|
|
|
identity_loader.clear("A")
|
|
|
|
a2, b2 = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a2 == "A"
|
|
assert b2 == "B"
|
|
|
|
assert load_calls == [["A", "B"], ["A"]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_clears_all_values_in_loader():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
a, b = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a == "A"
|
|
assert b == "B"
|
|
|
|
assert load_calls == [["A", "B"]]
|
|
|
|
identity_loader.clear_all()
|
|
|
|
a2, b2 = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a2 == "A"
|
|
assert b2 == "B"
|
|
|
|
assert load_calls == [["A", "B"], ["A", "B"]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_allows_priming_the_cache():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
identity_loader.prime("A", "A")
|
|
|
|
a, b = await gather(identity_loader.load("A"), identity_loader.load("B"))
|
|
|
|
assert a == "A"
|
|
assert b == "B"
|
|
|
|
assert load_calls == [["B"]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_does_not_prime_keys_that_already_exist():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
identity_loader.prime("A", "X")
|
|
|
|
a1 = await identity_loader.load("A")
|
|
b1 = await identity_loader.load("B")
|
|
|
|
assert a1 == "X"
|
|
assert b1 == "B"
|
|
|
|
identity_loader.prime("A", "Y")
|
|
identity_loader.prime("B", "Y")
|
|
|
|
a2 = await identity_loader.load("A")
|
|
b2 = await identity_loader.load("B")
|
|
|
|
assert a2 == "X"
|
|
assert b2 == "B"
|
|
|
|
assert load_calls == [["B"]]
|
|
|
|
|
|
# # Represents Errors
|
|
@mark.asyncio
|
|
async def test_resolves_to_error_to_indicate_failure():
|
|
async def resolve(keys):
|
|
mapped_keys = [
|
|
key if key % 2 == 0 else Exception("Odd: {}".format(key)) for key in keys
|
|
]
|
|
return mapped_keys
|
|
|
|
even_loader, load_calls = id_loader(resolve=resolve)
|
|
|
|
with raises(Exception) as exc_info:
|
|
await even_loader.load(1)
|
|
|
|
assert str(exc_info.value) == "Odd: 1"
|
|
|
|
value2 = await even_loader.load(2)
|
|
assert value2 == 2
|
|
assert load_calls == [[1], [2]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_can_represent_failures_and_successes_simultaneously():
|
|
async def resolve(keys):
|
|
mapped_keys = [
|
|
key if key % 2 == 0 else Exception("Odd: {}".format(key)) for key in keys
|
|
]
|
|
return mapped_keys
|
|
|
|
even_loader, load_calls = id_loader(resolve=resolve)
|
|
|
|
promise1 = even_loader.load(1)
|
|
promise2 = even_loader.load(2)
|
|
|
|
with raises(Exception) as exc_info:
|
|
await promise1
|
|
|
|
assert str(exc_info.value) == "Odd: 1"
|
|
value2 = await promise2
|
|
assert value2 == 2
|
|
assert load_calls == [[1, 2]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_caches_failed_fetches():
|
|
async def resolve(keys):
|
|
mapped_keys = [Exception("Error: {}".format(key)) for key in keys]
|
|
return mapped_keys
|
|
|
|
error_loader, load_calls = id_loader(resolve=resolve)
|
|
|
|
with raises(Exception) as exc_info:
|
|
await error_loader.load(1)
|
|
|
|
assert str(exc_info.value) == "Error: 1"
|
|
|
|
with raises(Exception) as exc_info:
|
|
await error_loader.load(1)
|
|
|
|
assert str(exc_info.value) == "Error: 1"
|
|
|
|
assert load_calls == [[1]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_caches_failed_fetches_2():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
identity_loader.prime(1, Exception("Error: 1"))
|
|
|
|
with raises(Exception) as _:
|
|
await identity_loader.load(1)
|
|
|
|
assert load_calls == []
|
|
|
|
|
|
# It is resilient to job queue ordering
|
|
@mark.asyncio
|
|
async def test_batches_loads_occuring_within_promises():
|
|
identity_loader, load_calls = id_loader()
|
|
|
|
async def load_b_1():
|
|
return await load_b_2()
|
|
|
|
async def load_b_2():
|
|
return await identity_loader.load("B")
|
|
|
|
values = await gather(identity_loader.load("A"), load_b_1())
|
|
|
|
assert values == ["A", "B"]
|
|
|
|
assert load_calls == [["A", "B"]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_catches_error_if_loader_resolver_fails():
|
|
exc = Exception("AOH!")
|
|
|
|
def do_resolve(x):
|
|
raise exc
|
|
|
|
a_loader, a_load_calls = id_loader(resolve=do_resolve)
|
|
|
|
with raises(Exception) as exc_info:
|
|
await a_loader.load("A1")
|
|
|
|
assert exc_info.value == exc
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_can_call_a_loader_from_a_loader():
|
|
deep_loader, deep_load_calls = id_loader()
|
|
a_loader, a_load_calls = id_loader(
|
|
resolve=lambda keys: deep_loader.load(tuple(keys))
|
|
)
|
|
b_loader, b_load_calls = id_loader(
|
|
resolve=lambda keys: deep_loader.load(tuple(keys))
|
|
)
|
|
|
|
a1, b1, a2, b2 = await gather(
|
|
a_loader.load("A1"),
|
|
b_loader.load("B1"),
|
|
a_loader.load("A2"),
|
|
b_loader.load("B2"),
|
|
)
|
|
|
|
assert a1 == "A1"
|
|
assert b1 == "B1"
|
|
assert a2 == "A2"
|
|
assert b2 == "B2"
|
|
|
|
assert a_load_calls == [["A1", "A2"]]
|
|
assert b_load_calls == [["B1", "B2"]]
|
|
assert deep_load_calls == [[("A1", "A2"), ("B1", "B2")]]
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_dataloader_clear_with_missing_key_works():
|
|
async def do_resolve(x):
|
|
return x
|
|
|
|
a_loader, a_load_calls = id_loader(resolve=do_resolve)
|
|
assert a_loader.clear("A1") == a_loader
|