Skip to content

Commit 3e43052

Browse files
committed
Add basic test for aiodataloader
1 parent a0b522f commit 3e43052

File tree

1 file changed

+79
-0
lines changed

1 file changed

+79
-0
lines changed

tests_asyncio/test_dataloader.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
from collections import namedtuple
2+
from unittest.mock import Mock
3+
from pytest import mark
4+
from aiodataloader import DataLoader
5+
6+
from graphene import ObjectType, String, Schema, Field, List
7+
8+
9+
CHARACTERS = {
10+
"1": {"name": "Luke Skywalker", "sibling": "3"},
11+
"2": {"name": "Darth Vader", "sibling": None},
12+
"3": {"name": "Leia Organa", "sibling": "1"},
13+
}
14+
15+
16+
get_character = Mock(side_effect=lambda character_id: CHARACTERS[character_id])
17+
18+
19+
class CharacterType(ObjectType):
20+
name = String()
21+
sibling = Field(lambda: CharacterType)
22+
23+
async def resolve_sibling(character, info):
24+
if character["sibling"]:
25+
return await info.context.character_loader.load(character["sibling"])
26+
return None
27+
28+
29+
class Query(ObjectType):
30+
skywalker_family = List(CharacterType)
31+
32+
async def resolve_skywalker_family(_, info):
33+
return await info.context.character_loader.load_many(["1", "2", "3"])
34+
35+
36+
mock_batch_load_fn = Mock(
37+
side_effect=lambda character_ids: [get_character(id) for id in character_ids]
38+
)
39+
40+
41+
class CharacterLoader(DataLoader):
42+
async def batch_load_fn(self, character_ids):
43+
return mock_batch_load_fn(character_ids)
44+
45+
46+
Context = namedtuple("Context", "character_loader")
47+
48+
49+
@mark.asyncio
50+
async def test_basic_dataloader():
51+
schema = Schema(query=Query)
52+
53+
character_loader = CharacterLoader()
54+
context = Context(character_loader=character_loader)
55+
56+
query = """
57+
{
58+
skywalkerFamily {
59+
name
60+
sibling {
61+
name
62+
}
63+
}
64+
}
65+
"""
66+
67+
result = await schema.execute_async(query, context=context)
68+
69+
assert not result.errors
70+
assert result.data == {
71+
"skywalkerFamily": [
72+
{"name": "Luke Skywalker", "sibling": {"name": "Leia Organa"}},
73+
{"name": "Darth Vader", "sibling": None},
74+
{"name": "Leia Organa", "sibling": {"name": "Luke Skywalker"}},
75+
]
76+
}
77+
78+
assert mock_batch_load_fn.call_count == 1
79+
assert get_character.call_count == 3

0 commit comments

Comments
 (0)