Fix tests to use proper ast node (#254)

- Update test_visitor.py to properly type-annotate the visitor class attribute
  and add assertion before using selection_set
- Update test_schema_parser.py to use more precise types that match GraphQL spec:
  - NonNullTypeNode's inner type can only be NamedTypeNode or ListTypeNode
  - Schema definitions use ConstDirectiveNode, not DirectiveNode
  - Default values use ConstValueNode, not ValueNode
  - OperationTypeDefinition's type_ must be NamedTypeNode

Author: corydolphin

tests/error/test_graphql_error.py

@@ -4,7 +4,7 @@
 
 from graphql.error import GraphQLError
 from graphql.language import (
-    Node,
+    NameNode,
     ObjectTypeDefinitionNode,
     OperationDefinitionNode,
     Source,
@@ -352,7 +352,7 @@ def formats_graphql_error():
         extensions = {"ext": None}
         error = GraphQLError(
             "test message",
-            Node(),
+            NameNode(value="stub"),
             Source(
                 """
                 query {

tests/language/test_ast.py

@@ -10,15 +10,22 @@
 class SampleTestNode(Node):
     __slots__ = "alpha", "beta"
 
-    alpha: int
-    beta: int
+    alpha: int | Node  # Union with Node to support copy tests with nested nodes
+    beta: int | Node | None
 
 
 class SampleNamedNode(Node):
     __slots__ = "foo", "name"
 
     foo: str
-    name: str | None
+    name: NameNode | None
+
+
+def make_loc(start: int = 1, end: int = 3) -> Location:
+    """Create a Location for testing with the given start/end offsets."""
+    source = Source("test source")
+    start_token = Token(TokenKind.NAME, start, end, 1, start, "test")
+    return Location(start_token, start_token, source)
 
 
 def describe_token_class():
@@ -150,15 +157,21 @@ def can_hash():
 
 def describe_node_class():
     def initializes_with_keywords():
-        node = SampleTestNode(alpha=1, beta=2, loc=0)
+        node = SampleTestNode(alpha=1, beta=2)
         assert node.alpha == 1
         assert node.beta == 2
-        assert node.loc == 0
-        node = SampleTestNode(alpha=1, loc=None)
         assert node.loc is None
+
+    def initializes_with_location():
+        loc = make_loc()
+        node = SampleTestNode(alpha=1, beta=2, loc=loc)
         assert node.alpha == 1
-        assert node.beta is None
-        node = SampleTestNode(alpha=1, beta=2, gamma=3)
+        assert node.beta == 2
+        assert node.loc is loc
+
+    def initializes_with_none_location():
+        node = SampleTestNode(alpha=1, beta=2, loc=None)
+        assert node.loc is None
         assert node.alpha == 1
         assert node.beta == 2
         assert not hasattr(node, "gamma")
@@ -174,27 +187,31 @@ def converts_list_to_tuple_on_init():
     def has_representation_with_loc():
         node = SampleTestNode(alpha=1, beta=2)
         assert repr(node) == "SampleTestNode"
-        node = SampleTestNode(alpha=1, beta=2, loc=3)
-        assert repr(node) == "SampleTestNode at 3"
+        loc = make_loc(start=3, end=5)
+        node = SampleTestNode(alpha=1, beta=2, loc=loc)
+        assert repr(node) == "SampleTestNode at 3:5"
 
     def has_representation_when_named():
         name_node = NameNode(value="baz")
         node = SampleNamedNode(foo="bar", name=name_node)
         assert repr(node) == "SampleNamedNode(name='baz')"
-        node = SampleNamedNode(alpha=1, beta=2, name=name_node, loc=3)
-        assert repr(node) == "SampleNamedNode(name='baz') at 3"
+        loc = make_loc(start=3, end=5)
+        node = SampleNamedNode(foo="bar", name=name_node, loc=loc)
+        assert repr(node) == "SampleNamedNode(name='baz') at 3:5"
 
     def has_representation_when_named_but_name_is_none():
-        node = SampleNamedNode(alpha=1, beta=2, name=None)
+        node = SampleNamedNode(foo="bar", name=None)
         assert repr(node) == "SampleNamedNode"
-        node = SampleNamedNode(alpha=1, beta=2, name=None, loc=3)
-        assert repr(node) == "SampleNamedNode at 3"
+        loc = make_loc(start=3, end=5)
+        node = SampleNamedNode(foo="bar", name=None, loc=loc)
+        assert repr(node) == "SampleNamedNode at 3:5"
 
     def has_special_representation_when_it_is_a_name_node():
         node = NameNode(value="foo")
         assert repr(node) == "NameNode('foo')"
-        node = NameNode(value="foo", loc=3)
-        assert repr(node) == "NameNode('foo') at 3"
+        loc = make_loc(start=3, end=5)
+        node = NameNode(value="foo", loc=loc)
+        assert repr(node) == "NameNode('foo') at 3:5"
 
     def can_check_equality():
         node = SampleTestNode(alpha=1, beta=2)

tests/language/test_schema_parser.py

@@ -3,14 +3,15 @@
 import pickle
 from copy import deepcopy
 from textwrap import dedent
-from typing import Optional, Tuple
 
 import pytest
 
 from graphql.error import GraphQLSyntaxError
 from graphql.language import (
     ArgumentNode,
     BooleanValueNode,
+    ConstDirectiveNode,
+    ConstValueNode,
     DirectiveDefinitionNode,
     DirectiveNode,
     DocumentNode,
@@ -22,6 +23,7 @@
     InterfaceTypeDefinitionNode,
     InterfaceTypeExtensionNode,
     ListTypeNode,
+    Location,
     NamedTypeNode,
     NameNode,
     NonNullTypeNode,
@@ -32,25 +34,30 @@
     ScalarTypeDefinitionNode,
     SchemaDefinitionNode,
     SchemaExtensionNode,
+    Source,
     StringValueNode,
+    Token,
+    TokenKind,
     TypeNode,
     UnionTypeDefinitionNode,
-    ValueNode,
     parse,
 )
 
 from ..fixtures import kitchen_sink_sdl  # noqa: F401
 
-try:
-    from typing import TypeAlias
-except ImportError:  # Python < 3.10
-    from typing_extensions import TypeAlias
 
-
-Location: TypeAlias = Optional[Tuple[int, int]]
+def make_loc(position: tuple[int, int]) -> Location:
+    """Create a Location for testing with the given (start, end) offsets."""
+    source = Source(body="")
+    token = Token(
+        kind=TokenKind.NAME, start=position[0], end=position[1], line=1, column=1
+    )
+    return Location(start_token=token, end_token=token, source=source)
 
 
-def assert_syntax_error(text: str, message: str, location: Location) -> None:
+def assert_syntax_error(
+    text: str, message: str, location: tuple[int, int] | None
+) -> None:
     with pytest.raises(GraphQLSyntaxError) as exc_info:
         parse(text)
     error = exc_info.value
@@ -59,85 +66,104 @@ def assert_syntax_error(text: str, message: str, location: Location) -> None:
     assert error.locations == [location]
 
 
-def assert_definitions(body: str, loc: Location, num=1):
+def assert_definitions(body: str, position: tuple[int, int] | None, num: int = 1):
     doc = parse(body)
     assert isinstance(doc, DocumentNode)
-    assert doc.loc == loc
+    assert doc.loc == position
     definitions = doc.definitions
     assert isinstance(definitions, tuple)
     assert len(definitions) == num
     return definitions[0] if num == 1 else definitions
 
 
-def type_node(name: str, loc: Location):
-    return NamedTypeNode(name=name_node(name, loc), loc=loc)
+def type_node(name: str, position: tuple[int, int]):
+    return NamedTypeNode(name=name_node(name, position), loc=make_loc(position))
 
 
-def name_node(name: str, loc: Location):
-    return NameNode(value=name, loc=loc)
+def name_node(name: str, position: tuple[int, int]):
+    return NameNode(value=name, loc=make_loc(position))
 
 
-def field_node(name: NameNode, type_: TypeNode, loc: Location):
-    return field_node_with_args(name, type_, (), loc)
+def field_node(name: NameNode, type_: TypeNode, position: tuple[int, int]):
+    return field_node_with_args(name, type_, (), position)
 
 
-def field_node_with_args(name: NameNode, type_: TypeNode, args: tuple, loc: Location):
+def field_node_with_args(
+    name: NameNode, type_: TypeNode, args: tuple, position: tuple[int, int]
+):
     return FieldDefinitionNode(
-        name=name, arguments=args, type=type_, directives=(), loc=loc, description=None
+        name=name,
+        arguments=args,
+        type=type_,
+        directives=(),
+        loc=make_loc(position),
+        description=None,
     )
 
 
-def non_null_type(type_: TypeNode, loc: Location):
-    return NonNullTypeNode(type=type_, loc=loc)
+def non_null_type(type_: NamedTypeNode | ListTypeNode, position: tuple[int, int]):
+    return NonNullTypeNode(type=type_, loc=make_loc(position))
 
 
-def enum_value_node(name: str, loc: Location):
+def enum_value_node(name: str, position: tuple[int, int]):
     return EnumValueDefinitionNode(
-        name=name_node(name, loc), directives=(), loc=loc, description=None
+        name=name_node(name, position),
+        directives=(),
+        loc=make_loc(position),
+        description=None,
     )
 
 
 def input_value_node(
-    name: NameNode, type_: TypeNode, default_value: ValueNode | None, loc: Location
+    name: NameNode,
+    type_: TypeNode,
+    default_value: ConstValueNode | None,
+    position: tuple[int, int],
 ):
     return InputValueDefinitionNode(
         name=name,
         type=type_,
         default_value=default_value,
         directives=(),
-        loc=loc,
+        loc=make_loc(position),
         description=None,
     )
 
 
-def boolean_value_node(value: bool, loc: Location):
-    return BooleanValueNode(value=value, loc=loc)
+def boolean_value_node(value: bool, position: tuple[int, int]):
+    return BooleanValueNode(value=value, loc=make_loc(position))
 
 
-def string_value_node(value: str, block: bool | None, loc: Location):
-    return StringValueNode(value=value, block=block, loc=loc)
+def string_value_node(value: str, block: bool | None, position: tuple[int, int]):
+    return StringValueNode(value=value, block=block, loc=make_loc(position))
 
 
-def list_type_node(type_: TypeNode, loc: Location):
-    return ListTypeNode(type=type_, loc=loc)
+def list_type_node(type_: TypeNode, position: tuple[int, int]):
+    return ListTypeNode(type=type_, loc=make_loc(position))
 
 
 def schema_extension_node(
-    directives: tuple[DirectiveNode, ...],
+    directives: tuple[ConstDirectiveNode, ...],
     operation_types: tuple[OperationTypeDefinitionNode, ...],
-    loc: Location,
+    position: tuple[int, int],
 ):
     return SchemaExtensionNode(
-        directives=directives, operation_types=operation_types, loc=loc
+        directives=directives, operation_types=operation_types, loc=make_loc(position)
     )
 
 
-def operation_type_definition(operation: OperationType, type_: TypeNode, loc: Location):
-    return OperationTypeDefinitionNode(operation=operation, type=type_, loc=loc)
+def operation_type_definition(
+    operation: OperationType, type_: NamedTypeNode, position: tuple[int, int]
+):
+    return OperationTypeDefinitionNode(
+        operation=operation, type=type_, loc=make_loc(position)
+    )
 
 
-def directive_node(name: NameNode, arguments: tuple[ArgumentNode, ...], loc: Location):
-    return DirectiveNode(name=name, arguments=arguments, loc=loc)
+def directive_node(
+    name: NameNode, arguments: tuple[ArgumentNode, ...], position: tuple[int, int]
+):
+    return DirectiveNode(name=name, arguments=arguments, loc=make_loc(position))
 
 
 def describe_schema_parser():

tests/language/test_visitor.py

@@ -308,7 +308,7 @@ def allows_editing_a_node_both_on_enter_and_on_leave():
         visited = []
 
         class TestVisitor(Visitor):
-            selection_set = None
+            selection_set: SelectionSetNode | None = None
 
             def enter_operation_definition(self, *args):
                 check_visitor_fn_args(ast, *args)
@@ -330,6 +330,7 @@ def leave_operation_definition(self, *args):
                 check_visitor_fn_args_edited(ast, *args)
                 node = args[0]
                 assert not node.selection_set.selections
+                assert self.selection_set is not None
                 # Create new node with original selection set (immutable pattern)
                 new_node = OperationDefinitionNode(
                     operation=node.operation,