adding @fromContext directive validation

apollo-federation/src/error/mod.rs

@@ -475,6 +475,8 @@ pub enum SingleFederationError {
     #[error("@cost cannot be applied to interface \"{interface}.{field}\"")]
     CostAppliedToInterfaceField { interface: Name, field: Name },
     #[error("{message}")]
+    ContextSelectionInvalid { message: String },
+    #[error("{message}")]
     ListSizeAppliedToNonList { message: String },
     #[error("{message}")]
     ListSizeInvalidAssumedSize { message: String },
@@ -689,6 +691,9 @@ impl SingleFederationError {
             SingleFederationError::ContextNoResolvableKey { .. } => {
                 ErrorCode::ContextNoResolvableKey
             }
+            SingleFederationError::ContextSelectionInvalid { .. } => {
+                ErrorCode::ContextSelectionInvalid
+            }
             SingleFederationError::CostAppliedToInterfaceField { .. } => {
                 ErrorCode::CostAppliedToInterfaceField
             }
@@ -1939,6 +1944,17 @@ static CONTEXT_NO_RESOLVABLE_KEY: LazyLock<ErrorCodeDefinition> = LazyLock::new(
     )
 });
 
+static CONTEXT_SELECTION_INVALID: LazyLock<ErrorCodeDefinition> = LazyLock::new(|| {
+    ErrorCodeDefinition::new(
+        "CONTEXT_SELECTION_INVALID".to_owned(),
+        "The selection set is invalid".to_owned(),
+        Some(ErrorCodeMetadata {
+            added_in: "2.8.0",
+            replaces: &[],
+        }),
+    )
+});
+
 #[derive(Debug, PartialEq, strum_macros::EnumIter)]
 pub enum ErrorCode {
     Internal,
@@ -2033,6 +2049,7 @@ pub enum ErrorCode {
     NoContextReferenced,
     NoSelectionForContext,
     ContextNoResolvableKey,
+    ContextSelectionInvalid,
 }
 
 impl ErrorCode {
@@ -2145,6 +2162,7 @@ impl ErrorCode {
             ErrorCode::NoContextReferenced => &NO_CONTEXT_REFERENCED,
             ErrorCode::NoSelectionForContext => &NO_SELECTION_FOR_CONTEXT,
             ErrorCode::ContextNoResolvableKey => &CONTEXT_NO_RESOLVABLE_KEY,
+            ErrorCode::ContextSelectionInvalid => &CONTEXT_SELECTION_INVALID,
         }
     }
 }

apollo-federation/src/schema/schema_upgrader.rs

@@ -490,7 +490,7 @@ impl SchemaUpgrader {
                         .federation_spec_definition()
                         .key_directive_arguments(directive)?;
                     for field in collect_target_fields_from_field_set(
-                        Valid::assume_valid_ref(schema.schema()),
+                        Valid::assume_valid_ref(other_schema.schema().schema()),
                         ty.type_name().clone(),
                         args.fields,
                         false,

apollo-federation/src/schema/validators/context.rs

@@ -100,7 +100,6 @@ impl ContextValidator for DenyInvalidContextName {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::subgraph::test_utils::build_and_expand;
 
     #[test]
     fn deny_underscore_in_context_name() {
@@ -162,59 +161,4 @@ mod tests {
         rule.validate("validName123", &mut errors);
         assert_eq!(errors.errors.len(), 0, "Expected no errors for valid name");
     }
-
-    #[test]
-    fn builds_context_type_map() {
-        let schema_str = r#"
-            extend schema
-                @link(url: "https://specs.apollo.dev/federation/v2.8", import: ["@context"])
-
-            type Query {
-                t: T
-            }
-
-            type T @context(name: "contextA") {
-                id: ID!
-            }
-
-            type U @context(name: "contextA") {
-                id: ID!
-            }
-
-            type V @context(name: "contextB") {
-                id: ID!
-            }
-        "#;
-
-        let subgraph = build_and_expand(schema_str);
-        let mut errors = MultipleFederationErrors::new();
-        let context_map =
-            validate_context_directives(subgraph.schema(), &mut errors).expect("validates");
-
-        // Check that there are no validation errors
-        assert_eq!(errors.errors.len(), 0, "Expected no validation errors");
-
-        // Check the context map contents
-        assert_eq!(context_map.len(), 2, "Expected two context names");
-
-        // Check contextA types
-        let context_a_types = context_map.get("contextA").expect("contextA should exist");
-        assert_eq!(context_a_types.len(), 2, "contextA should have two types");
-        assert!(
-            context_a_types.iter().any(|name| name == "T"),
-            "contextA should include type T"
-        );
-        assert!(
-            context_a_types.iter().any(|name| name == "U"),
-            "contextA should include type U"
-        );
-
-        // Check contextB types
-        let context_b_types = context_map.get("contextB").expect("contextB should exist");
-        assert_eq!(context_b_types.len(), 1, "contextB should have one type");
-        assert!(
-            context_b_types.iter().any(|name| name == "V"),
-            "contextB should include type V"
-        );
-    }
 }