From cd216447c256f4aef8501653955ea2d80eb993a0 Mon Sep 17 00:00:00 2001
From: Syrus Akbary <me@syrusakbary.com>
Date: Fri, 25 Sep 2015 16:36:18 -0700
Subject: [PATCH] Added test relay schema

---
 tests/starwars_relay/schema.py | 36 +++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/tests/starwars_relay/schema.py b/tests/starwars_relay/schema.py
index 1a7e570d..a6a0a750 100644
--- a/tests/starwars_relay/schema.py
+++ b/tests/starwars_relay/schema.py
@@ -1,47 +1,49 @@
 import graphene
-from graphene import resolve_only_args
+from graphene import resolve_only_args, relay
 
-from .data import getHero, getHuman, getCharacter, getDroid, Human as _Human, Droid as _Droid
+from .data import (
+    getHero, getHuman, getCharacter, getDroid,
+    Human as _Human, Droid as _Droid)
 
 Episode = graphene.Enum('Episode', dict(
-    NEWHOPE = 4,
-    EMPIRE = 5,
-    JEDI = 6
+    NEWHOPE=4,
+    EMPIRE=5,
+    JEDI=6
 ))
 
+
 def wrap_character(character):
-    if isinstance(character, _Human): 
+    if isinstance(character, _Human):
         return Human(character)
     elif isinstance(character, _Droid):
         return Droid(character)
 
+
 class Character(graphene.Interface):
-    id = graphene.IDField()
     name = graphene.StringField()
-    friends = graphene.ListField('self')
+    friends = relay.Connection('self')
     appearsIn = graphene.ListField(Episode)
 
     def resolve_friends(self, args, *_):
         return [wrap_character(getCharacter(f)) for f in self.instance.friends]
 
-class Human(Character):
+
+class Human(relay.Node, Character):
     homePlanet = graphene.StringField()
 
 
-class Droid(Character):
+class Droid(relay.Node, Character):
     primaryFunction = graphene.StringField()
 
 
 class Query(graphene.ObjectType):
     hero = graphene.Field(Character,
-        episode = graphene.Argument(Episode)
-    )
+                          episode=graphene.Argument(Episode))
     human = graphene.Field(Human,
-        id = graphene.Argument(graphene.String)
-    )
+                           id=graphene.Argument(graphene.String))
     droid = graphene.Field(Droid,
-        id = graphene.Argument(graphene.String)
-    )
+                           id=graphene.Argument(graphene.String))
+    node = graphene.Field(relay.Node)
 
     @resolve_only_args
     def resolve_hero(self, episode):
@@ -50,8 +52,6 @@ class Query(graphene.ObjectType):
     @resolve_only_args
     def resolve_human(self, id):
         return wrap_character(getHuman(id))
-        if human:
-            return Human(human)
 
     @resolve_only_args
     def resolve_droid(self, id):