apollo-language-server 0.6.0

use crate::{
    completion::get_completions_at_position,
    federation::{
        get_connect_builtins, get_federation_builtins_for_v1_document,
        get_federation_builtins_for_v2_document, get_missing_query_type_for_document,
        get_tag_builtins,
        link::{is_link_compatible, ParsedLink},
        SpecBuiltins,
    },
    semantic_tokens::IncompleteSemanticToken,
    server::MaxSpecVersions,
    specs::{
        federation::{
            EXTENDS_DIRECTIVE, EXTERNAL_DIRECTIVE, KEY_DIRECTIVE, PROVIDES_DIRECTIVE,
            REQUIRES_DIRECTIVE,
        },
        link::LINK_DIRECTIVE,
    },
};
use apollo_compiler::{
    ast::{Definition, Document},
    parser::Parser,
    validation::DiagnosticList,
    Schema,
};
use monolith::Monolith;
use std::collections::HashMap;
use subgraph::{fed1_definitions_are_compatible, Subgraph, SubgraphFederationVersion};
use supergraph::{KnownSubgraphs, Supergraph};

pub mod monolith;
pub mod subgraph;
pub mod supergraph;

#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum Graph {
    Monolith(monolith::Monolith),
    Supergraph(Box<supergraph::Supergraph>),
}

#[derive(Default, Clone)]
pub struct GraphConfig {
    pub force_federation: bool,
    /// The dynamically fetched federation spec to use instead of the built-in copies
    pub federation_spec: Option<String>,
}

impl Graph {
    pub fn new(
        uri: lsp::Url,
        source_text: String,
        version: i32,
        known_subgraphs: KnownSubgraphs,
        config: GraphConfig,
    ) -> Graph {
        let schema_with_metadata =
            parse_as_subgraph_or_monolith(&source_text, uri.as_str(), config);

        if schema_with_metadata.is_federated {
            let mut supergraph = Supergraph::new(known_subgraphs);

            let name = uri
                .query_pairs()
                .find_map(|(key, value)| (key == "subgraph_name").then(|| value.to_string()))
                .unwrap_or_else(|| {
                    supergraph
                        .known_subgraphs
                        .by_uri
                        .get(&uri)
                        .cloned()
                        .unwrap_or_else(|| uri.to_string())
                });

            supergraph.insert(Subgraph::new(
                name,
                uri,
                source_text,
                version,
                schema_with_metadata,
            ));
            Graph::Supergraph(Box::new(supergraph))
        } else {
            Graph::Monolith(Monolith::new(
                uri,
                source_text,
                version,
                schema_with_metadata,
            ))
        }
    }

    pub fn update(
        &mut self,
        uri: lsp::Url,
        source_text: String,
        version: i32,
        config: GraphConfig,
    ) {
        let schema_with_metadata =
            parse_as_subgraph_or_monolith(&source_text, uri.as_str(), config);
        match self {
            Graph::Monolith(monolith) => {
                monolith.update(source_text, version, schema_with_metadata)
            }
            Graph::Supergraph(supergraph) => {
                let name = uri
                    .query_pairs()
                    .find_map(|(key, value)| (key == "subgraph_name").then(|| value.to_string()))
                    .unwrap_or_else(|| {
                        supergraph
                            .known_subgraphs
                            .by_uri
                            .get(&uri)
                            .cloned()
                            .unwrap_or_else(|| uri.to_string())
                    });

                supergraph.update(Subgraph::new(
                    name,
                    uri,
                    source_text,
                    version,
                    schema_with_metadata,
                ))
            }
        }
    }

    pub fn diagnostics_for_uri(&self, uri: &lsp::Url) -> (Vec<lsp::Diagnostic>, i32) {
        match self {
            Graph::Monolith(monolith) => (monolith.diagnostics(), monolith.version),
            Graph::Supergraph(supergraph) => supergraph.diagnostics_for_subgraph(uri),
        }
    }

    pub(crate) fn version_for_uri(&self, uri: &lsp::Url) -> Option<i32> {
        match self {
            Graph::Monolith(monolith) => Some(monolith.version),
            Graph::Supergraph(supergraph) => supergraph.version_for_subgraph(uri),
        }
    }

    pub fn completions(
        &self,
        subgraph_uri: &lsp::Url,
        position: lsp::Position,
        max_spec_versions: &MaxSpecVersions,
    ) -> Option<lsp::CompletionResponse> {
        get_completions_at_position(self, subgraph_uri, position, max_spec_versions)
    }

    pub fn semantic_tokens_full(
        &self,
        subgraph_uri: &lsp::Url,
    ) -> Option<Vec<IncompleteSemanticToken>> {
        match self {
            Graph::Monolith(_) => None,
            Graph::Supergraph(supergraph) => Some(
                supergraph
                    .subgraph_by_uri(subgraph_uri)
                    .unwrap()
                    .semantic_tokens(),
            ),
        }
    }

    pub fn on_hover(&self, uri: &lsp::Url, position: &lsp::Position) -> Option<lsp::Hover> {
        match self {
            Graph::Monolith(monolith) => monolith.on_hover(position),
            Graph::Supergraph(supergraph) => supergraph.on_hover(uri, position),
        }
    }

    pub fn goto_definition(
        &self,
        uri: &lsp::Url,
        position: &lsp::Position,
    ) -> Option<Vec<lsp::LocationLink>> {
        match self {
            Graph::Monolith(monolith) => monolith.goto_definition(position),
            Graph::Supergraph(supergraph) => supergraph.goto_definition(uri, position),
        }
    }

    pub fn supergraph(&self) -> Option<&supergraph::Supergraph> {
        match self {
            Graph::Supergraph(supergraph) => Some(supergraph),
            _ => None,
        }
    }

    #[cfg(test)]
    pub fn source_text_for_uri(&self, uri: &lsp::Url) -> &ropey::Rope {
        match self {
            Graph::Monolith(monolith) => &monolith.source_text,
            Graph::Supergraph(supergraph) => &supergraph.subgraph_by_uri(uri).unwrap().source_text,
        }
    }

    // #[cfg(test)]
    // pub fn monolith(&self) -> &monolith::Monolith {
    //     match self {
    //         Graph::Monolith(monolith) => monolith,
    //         _ => panic!("Expected Monolith"),
    //     }
    // }
}

#[derive(Debug)]
pub struct SchemaWithMetadata {
    pub is_federated: bool,
    pub parse_errors: Option<DiagnosticList>,
    pub build_errors: Option<DiagnosticList>,
    pub validation_errors: Option<DiagnosticList>,
    pub schema: Schema,
    pub specs_with_aliases: HashMap<String, String>,
    pub links: HashMap<String, ParsedLink>,
}
// This function is a wrapper around the apollo_compiler's `Schema::parse` and
// `Schema::validate`. Apollo-flavored subgraphs are technically invalid, and
// these validation errors will surface in the editor. In order to accommodate
// this, we need to inject the appropriate federation builtins into the schema.
pub fn parse_as_subgraph_or_monolith(
    source_text: &str,
    uri: &str,
    config: GraphConfig,
) -> SchemaWithMetadata {
    let (parser, document, parse_errors) = parse_unknown_graph(source_text, uri);

    let Some(federation_version) = detect_federation_version_in_document(&document).or(
        // If the user has `force_federation` enabled, default to fed V2 and build as a subgraph
        config
            .force_federation
            .then_some(SubgraphFederationVersion::V2),
    ) else {
        return build_monolith(document, parse_errors);
    };

    build_subgraph(
        document,
        federation_version,
        parser,
        config.federation_spec,
        parse_errors,
    )
}

pub fn parse_unknown_graph(
    source_text: &str,
    uri: &str,
) -> (Parser, Document, Option<DiagnosticList>) {
    let mut parser = Parser::new();

    let (document, parse_errors) = match parser.parse_ast(source_text, uri) {
        Ok(document) => (document, None),
        Err(document_with_errors) => (
            document_with_errors.partial,
            Some(document_with_errors.errors),
        ),
    };
    (parser, document, parse_errors)
}

pub fn build_subgraph(
    document: Document,
    federation_version: SubgraphFederationVersion,
    mut parser: Parser,
    federation_spec: Option<String>,
    parse_errors: Option<DiagnosticList>,
) -> SchemaWithMetadata {
    let mut builder = Schema::builder().adopt_orphan_extensions();
    builder = builder.add_ast(&document);

    let mut links = HashMap::new();
    let mut specs_with_aliases = HashMap::new();

    if let Some(dynamic_spec) = federation_spec {
        parser.parse_into_schema_builder(dynamic_spec, "federation_builtins.graphql", &mut builder);
    } else {
        let mut parse_into_schema_builder =
            |path: &str, get_builtins: &dyn Fn(&Document) -> SpecBuiltins| {
                if let Some((builtins, spec_names, parsed_link)) = get_builtins(&document) {
                    parser.parse_into_schema_builder(builtins, path, &mut builder);
                    if let Some(spec_names) = spec_names {
                        specs_with_aliases.extend(spec_names);
                    }
                    if let Some(parsed_link) = parsed_link {
                        links.insert(parsed_link.spec_name.clone(), parsed_link);
                    }
                }
            };

        parse_into_schema_builder("federation_builtins.graphql", &|document| {
            match federation_version {
                SubgraphFederationVersion::V1 => get_federation_builtins_for_v1_document(document),
                SubgraphFederationVersion::V2 => get_federation_builtins_for_v2_document(document),
            }
        });

        parse_into_schema_builder("connect_builtins.graphql", &get_connect_builtins);

        parse_into_schema_builder("tags_builtins.graphql", &get_tag_builtins);

        parse_into_schema_builder("missing_query_type.graphql", &|document| {
            get_missing_query_type_for_document(document).map(|builtins| (builtins, None, None))
        });
    }

    let (schema, schema_build_errors) = match builder.build() {
        Ok(schema) => (schema, None),
        Err(with_errors) => (with_errors.partial, Some(with_errors.errors)),
    };

    let (schema, schema_validation_errors) = match schema.validate() {
        Ok(valid_schema) => (valid_schema.into_inner(), None),
        Err(with_errors) => (with_errors.partial, Some(with_errors.errors)),
    };

    SchemaWithMetadata {
        is_federated: true,
        parse_errors,
        build_errors: schema_build_errors,
        validation_errors: schema_validation_errors,
        schema,
        specs_with_aliases,
        links,
    }
}

pub fn build_monolith(
    document: Document,
    parse_errors: Option<DiagnosticList>,
) -> SchemaWithMetadata {
    let (schema, build_errors) = match document.to_schema() {
        Ok(schema) => (schema, None),
        Err(with_errors) => (with_errors.partial, Some(with_errors.errors)),
    };

    let (schema, validation_errors) = match schema.validate() {
        Ok(valid_schema) => (valid_schema.into_inner(), None),
        Err(with_errors) => (with_errors.partial, Some(with_errors.errors)),
    };

    SchemaWithMetadata {
        is_federated: false,
        parse_errors,
        build_errors,
        validation_errors,
        schema,
        specs_with_aliases: HashMap::new(),
        links: HashMap::new(),
    }
}

pub fn detect_federation_version_in_document(
    document: &Document,
) -> Option<SubgraphFederationVersion> {
    let has_link_usage = document.definitions.iter().any(|def| {
        // collect the directives from the definition if it's a schema definition or extension
        def.as_schema_definition()
            .map(|def| &def.directives)
            .or_else(|| def.as_schema_extension().map(|ext| &ext.directives))
            // check if any of the directives are a `@link`
            .map(|schema_directives| {
                schema_directives
                    .iter()
                    .any(|directive| directive.name == LINK_DIRECTIVE)
            })
            .unwrap_or_default()
    });

    let directive_definitions = document
        .definitions
        .iter()
        .filter(|def| def.as_directive_definition().is_some())
        .collect::<Vec<_>>();

    let maybe_link_definition = directive_definitions
        .iter()
        .find(|def| def.name().unwrap() == LINK_DIRECTIVE);

    // Case 1: A `@link` usage with no `@link definition`
    if has_link_usage && maybe_link_definition.is_none() {
        return Some(SubgraphFederationVersion::V2);
    }

    // Case 2: A `@link` definition that's compatible with ours
    if let Some(def) = maybe_link_definition.map(|def| def.as_directive_definition().unwrap()) {
        if is_link_compatible(def) {
            return Some(SubgraphFederationVersion::V2);
        }
    }

    // Case 3: Any directive definition that's compatible with `@link`
    if directive_definitions
        .iter()
        .any(|def| is_link_compatible(def.as_directive_definition().unwrap()))
    {
        return Some(SubgraphFederationVersion::V2);
    }

    // Case 4: Fed 1 directive definitions that match ours
    let fed1_directive_definitions = directive_definitions
        .iter()
        .filter(|def| {
            matches!(
                def.name().unwrap().as_str(),
                KEY_DIRECTIVE
                    | REQUIRES_DIRECTIVE
                    | PROVIDES_DIRECTIVE
                    | EXTERNAL_DIRECTIVE
                    | EXTENDS_DIRECTIVE
            )
        })
        .map(|def| def.as_directive_definition().unwrap())
        .collect::<Vec<_>>();

    if !fed1_directive_definitions.is_empty() {
        // It's ok to bail early here if they're _incompatible_ as well
        return fed1_definitions_are_compatible(fed1_directive_definitions)
            .then_some(SubgraphFederationVersion::V1);
    }

    // Case 5: Fed 1 directive usages without corresponding definitions We
    // really only need to look for a single `@key` usage on a type or an
    // interface. None of the other directives make sense without a `@key`
    // usage on the parent type.
    let has_key_usages = document.definitions.iter().any(|def| {
        let directives = match def {
            Definition::ObjectTypeDefinition(ty) => &ty.directives,
            Definition::ObjectTypeExtension(ty) => &ty.directives,
            Definition::InterfaceTypeDefinition(ty) => &ty.directives,
            Definition::InterfaceTypeExtension(ty) => &ty.directives,
            _ => return false,
        };
        directives
            .iter()
            .any(|directive| directive.name == KEY_DIRECTIVE)
    });

    if has_key_usages {
        return Some(SubgraphFederationVersion::V1);
    }

    None
}

#[cfg(test)]
mod monolith_tests {
    use super::*;

    #[test]
    fn test_simple() {
        let source_text = r#"
            type Query {
                hello: String!
            }
        "#;
        let Graph::Monolith(monolith) = Graph::new(
            lsp::Url::parse("file:///test.graphql").unwrap(),
            source_text.to_string(),
            0,
            KnownSubgraphs::default(),
            GraphConfig::default(),
        ) else {
            panic!("Expected Monolith");
        };

        assert!(monolith.diagnostics().is_empty());
    }

    #[test]
    fn test_with_schema_definition() {
        let source_text = r#"
            schema {
                query: Query
            }

            type Query {
                hello: String!
            }
        "#;
        let Graph::Monolith(monolith) = Graph::new(
            lsp::Url::parse("file:///test.graphql").unwrap(),
            source_text.to_string(),
            0,
            KnownSubgraphs::default(),
            GraphConfig::default(),
        ) else {
            panic!("Expected Monolith");
        };

        assert!(monolith.diagnostics().is_empty());
    }
}

#[cfg(test)]
fn subgraph_uri() -> lsp::Url {
    lsp::Url::parse("file:///test.graphql").unwrap()
}

#[cfg(test)]
fn get_testing_supergraph(source_text: &str, config: Option<GraphConfig>) -> Box<Supergraph> {
    let config = config.unwrap_or_default();
    let Graph::Supergraph(supergraph) = Graph::new(
        subgraph_uri(),
        source_text.to_string(),
        0,
        KnownSubgraphs::default(),
        config,
    ) else {
        panic!("Expected Supergraph");
    };
    supergraph
}

#[cfg(test)]
mod fed1_subgraph_tests {
    use insta::assert_snapshot;

    use super::*;

    #[test]
    fn test_simple() {
        let source_text = r#"
            type Query {
                user: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#;
        let supergraph = get_testing_supergraph(source_text, None);

        assert!(supergraph
            .diagnostics_for_subgraph(&subgraph_uri())
            .0
            .is_empty());
    }

    #[test]
    fn test_parse_v1_subgraph_mixed() {
        let source_text = r#"
            directive @external on FIELD_DEFINITION

            directive @requires(fields: _FieldSet!) on FIELD_DEFINITION

            directive @provides(fields: _FieldSet!) on FIELD_DEFINITION

            directive @key(fields: _FieldSet!) on OBJECT | INTERFACE

            type SomeType @extends {
                someField: String
            }
        "#;

        let schema_with_metadata = parse_as_subgraph_or_monolith(
            source_text,
            "file:///test.graphql",
            GraphConfig {
                force_federation: true,
                federation_spec: None,
            },
        );

        assert!(schema_with_metadata.is_federated);
        assert!(schema_with_metadata.parse_errors.is_none());
        assert!(schema_with_metadata.build_errors.is_none());
        assert!(schema_with_metadata.validation_errors.is_none());
        assert_snapshot!(schema_with_metadata.schema);
    }

    #[test]
    fn test_infers_v1_subgraph_with_contact_directive() {
        let source_text = r#"
            directive @contact(
                name: String!
                url: String!
                description: String
            ) on SCHEMA

            extend schema @contact(
                name: "testing",
                url: "https://testing.com",
                description: "A testing subgraph"
            )

            type Query {
                product: Product!
            }

            type Product @key(fields: "id") {
                id: ID!
            }
        "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();

        let diagnostics = subgraph.diagnostics();
        assert!(diagnostics.is_empty());
        assert_snapshot!(subgraph.schema());
    }

    #[test]
    fn test_no_query_type_entity_any_defined() {
        let source_text = r#"
            scalar _Any

            union _Entity = MyEntity

            type MyEntity @key(fields: "id") {
                id: ID!
            }
        "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();

        let diagnostics = subgraph.diagnostics();
        assert!(diagnostics.is_empty());
        assert_snapshot!(subgraph.schema());
    }
}

#[cfg(test)]
mod fed2_subgraph_tests {
    use crate::specs::federation::OVERRIDE_DIRECTIVE;

    use super::*;
    use insta::{assert_debug_snapshot, assert_snapshot};

    fn with_v2_imports(source_text: &str) -> String {
        format!(
            "{}\n\n{}",
            r#"extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@requires", "@provides", "@extends", "@external", "@tag", "@override", "@shareable", "@inaccessible", "FieldSet"])"#,
            source_text
        )
    }

    #[test]
    fn test_simple() {
        let source_text = &with_v2_imports(
            r#"
            type Query {
                user: User
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#,
        );

        let supergraph = get_testing_supergraph(source_text, None);
        assert!(supergraph
            .diagnostics_for_subgraph(&subgraph_uri())
            .0
            .is_empty());
    }

    #[test]
    fn test_query_type_provided_via_schema_definition() {
        let source_text = r#"
            schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@requires", "@provides", "@extends", "@external", "@tag", "@override", "@shareable", "@inaccessible", "FieldSet"]) {
                query: Query
            }

            type Query {
                user: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        assert!(subgraph.diagnostics().is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_query_type_provided_via_schema_extension() {
        let source_text = &with_v2_imports(
            r#"
            type Query {
                user: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#,
        );

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        assert!(subgraph.diagnostics().is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_query_type_provided_with_implicit_schema_definition() {
        let source_text = &with_v2_imports(
            r#"
            type Query {
                user: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#,
        );

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        assert!(subgraph.diagnostics().is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_query_and_mutation_type_provided_with_implicit_schema_definition() {
        let source_text = &with_v2_imports(
            r#"
            type Query {
                user: User!
            }

            type Mutation {
                createUser: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#,
        );

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();

        assert!(subgraph.diagnostics().is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_no_query_type_provided() {
        let source_text = &with_v2_imports(
            r#"
            type User @key(fields: "id") {
                id: ID!
            }
        "#,
        );

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        assert!(subgraph.diagnostics().is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_force_federation_config() {
        let source_text = r#"
            extend schema @contact(
                description: "Send urgent issues to @{your team name}",
                name : "{Your Team}",
                url : "{your team's slack channel}"
            )

            directive @contact(
                "Contact title of the subgraph owner"
                name: String!,
                "Other relevant notes can be included here; supports markdown"
                description: String,
                "URL where the subgraph can be reached"
                url: String
            ) on SCHEMA
        "#;

        let supergraph = get_testing_supergraph(
            source_text,
            Some(GraphConfig {
                force_federation: true,
                federation_spec: None,
            }),
        );

        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let diagnostics = subgraph.diagnostics();
        assert!(diagnostics.is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_force_federation_config_correctly_infers_v1_subgraph() {
        let source_text = r#"
            type Query {
                user: User!
            }

            type User @key(fields: "id") {
                id: ID!
            }
        "#;

        let supergraph = get_testing_supergraph(
            source_text,
            Some(GraphConfig {
                force_federation: true,
                federation_spec: None,
            }),
        );

        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let diagnostics = subgraph.diagnostics();
        assert!(diagnostics.is_empty());
        assert_snapshot!(subgraph.schema().to_string());
    }

    #[test]
    fn test_semantic_highlights_for_key_directive() {
        let source_text = r#"
             extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key"])

            type Query {
                user: User!
            }

            type Mutation {
                createUser: User!
            }

            type Name @key(fields: "first") {
                first: String!
                last: String!
            }

            type User @key(fields: """
                                    id 
                                    name { first }
                                """) {
                id: ID!
                name: Name!
            }
        "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let tokens = subgraph.semantic_tokens();

        assert_eq!(tokens.len(), 6);
        assert_debug_snapshot!(tokens);
    }

    #[test]
    fn test_semantic_highlights_for_fieldsets() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@provides", "@requires"])

            type Query {
                user: User!
            }

            type Mutation {
                createUser: User!
            }

            type Name @key(fields: "first") {
                first: String! @requires(fields: "last")
                last: String! @external
            }

            type User {
                id: ID!
                name: Name! @provides(fields: "name")
            }
        "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let tokens = subgraph.semantic_tokens();

        assert_eq!(tokens.len(), 3);
        assert_debug_snapshot!(tokens);
    }

    #[test]
    fn test_semantic_highlights_for_fieldset_with_args() {
        let source_text = r#"
        extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@provides", "@requires"])

        type Query {
            user: User!
        }

        type Mutation {
            createUser: User!
        }

        type Name {
            first: String!
            last: String! @external
        }

        type User {
            id: ID!
            name: Name! @provides(fields: "name")
        }

        type Team @key(fields: "id") {
            id: ID!
            name: String!
            relatedTeam(name: String, value: Int): [Team!]!
            teamTag: String! @requires(fields: "relatedTeams(name:\"abc\", value: 123) { name }")
        }
    "#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let tokens = subgraph.semantic_tokens();

        assert_eq!(tokens.len(), 16);
        assert_debug_snapshot!(tokens);
    }

    #[test]
    fn test_spec_with_alias() {
        let source_text = r#"
        extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: [{ name: "@key", as: "@keyAlias" }, "@override"], as: "testing")
"#;

        let supergraph = get_testing_supergraph(source_text, None);
        let subgraph = supergraph.subgraph_by_uri(&subgraph_uri()).unwrap();
        let schema_with_metadata = subgraph.schema_with_metadata();
        assert!(schema_with_metadata.specs_with_aliases.len() > 1);
        assert_eq!(
            schema_with_metadata.specs_with_aliases.get(KEY_DIRECTIVE),
            Some(&"keyAlias".to_string())
        );
        assert_eq!(
            schema_with_metadata
                .specs_with_aliases
                .get(OVERRIDE_DIRECTIVE),
            Some(&"override".to_string())
        );
        assert_eq!(
            schema_with_metadata
                .specs_with_aliases
                .get(EXTERNAL_DIRECTIVE),
            Some(&"testing__external".to_string())
        );
    }

    #[test]
    fn test_tolerant_to_bad_imports() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key"-])

            type Query {
                hello: String!
            }
        "#;

        get_testing_supergraph(source_text, None);
    }
}