apollo_smith/
response.rs

1use apollo_compiler::executable::Selection;
2use apollo_compiler::executable::SelectionSet;
3use apollo_compiler::schema::ExtendedType;
4use apollo_compiler::validation::Valid;
5use apollo_compiler::ExecutableDocument;
6use apollo_compiler::Name;
7use apollo_compiler::Schema;
8use arbitrary::Result;
9use arbitrary::Unstructured;
10use serde_json_bytes::json;
11use serde_json_bytes::serde_json::Number;
12use serde_json_bytes::Map;
13use serde_json_bytes::Value;
14use std::collections::HashMap;
15
16const TYPENAME: &str = "__typename";
17
18pub type Generator = Box<dyn Fn(&mut Unstructured) -> Result<Value>>;
19
20/// Builds a GraphQL response which matches the shape of a given executable GraphQL document.
21///
22/// Documentation of the response format can be found in the [GraphQL spec](https://spec.graphql.org/draft/#sec-Execution-Result).
23pub struct ResponseBuilder<'a, 'doc, 'schema> {
24    u: &'a mut Unstructured<'a>,
25    doc: &'doc Valid<ExecutableDocument>,
26    schema: &'schema Valid<Schema>,
27    custom_scalar_generators: HashMap<Name, Generator>,
28    min_list_size: usize,
29    max_list_size: usize,
30    null_ratio: Option<(u8, u8)>,
31    operation_name: Option<&'doc str>,
32}
33
34impl<'a, 'doc, 'schema> ResponseBuilder<'a, 'doc, 'schema> {
35    pub fn new(
36        u: &'a mut Unstructured<'a>,
37        doc: &'doc Valid<ExecutableDocument>,
38        schema: &'schema Valid<Schema>,
39    ) -> Self {
40        Self {
41            u,
42            doc,
43            schema,
44            custom_scalar_generators: HashMap::new(),
45            min_list_size: 0,
46            max_list_size: 5,
47            null_ratio: None,
48            operation_name: None,
49        }
50    }
51
52    /// Register a generator function for generating custom scalar values.
53    pub fn with_custom_scalar(mut self, scalar_name: Name, generator: Generator) -> Self {
54        self.custom_scalar_generators.insert(scalar_name, generator);
55        self
56    }
57
58    /// Set the minimum number of items per list field. If unset, defaults to 0.
59    pub fn with_min_list_size(mut self, min_size: usize) -> Self {
60        self.min_list_size = min_size;
61        self
62    }
63
64    /// Set the maximum number of items per list field. If unset, defaults to 5.
65    pub fn with_max_list_size(mut self, max_size: usize) -> Self {
66        self.max_list_size = max_size;
67        self
68    }
69
70    /// Set the frequency of null values for nullable fields. If unset, fields will never be null.
71    pub fn with_null_ratio(mut self, numerator: u8, denominator: u8) -> Self {
72        self.null_ratio = Some((numerator, denominator));
73        self
74    }
75
76    /// Set the operation name to generate a response for. If unset, uses the anonymous operation.
77    /// If the operation does not exist, returns a response with `data: null`.
78    pub fn with_operation_name(mut self, operation_name: Option<&'doc str>) -> Self {
79        self.operation_name = operation_name;
80        self
81    }
82
83    /// Builds a `Value` matching the shape of `self.doc`
84    pub fn build(mut self) -> Result<Value> {
85        if let Ok(operation) = self.doc.operations.get(self.operation_name) {
86            Ok(json!({ "data": self.arbitrary_selection_set(&operation.selection_set)? }))
87        } else {
88            Ok(json!({ "data": null }))
89        }
90    }
91
92    fn arbitrary_selection_set(&mut self, selection_set: &SelectionSet) -> Result<Value> {
93        let mut result = Map::new();
94
95        for selection in &selection_set.selections {
96            match selection {
97                Selection::Field(field) => {
98                    if field.name == TYPENAME {
99                        result.insert(
100                            field.name.to_string(),
101                            Value::String(selection_set.ty.to_string().into()),
102                        );
103                    } else if !field.ty().is_non_null() && self.should_be_null()? {
104                        result.insert(field.name.to_string(), Value::Null);
105                    } else if field.selection_set.is_empty() && !field.ty().is_list() {
106                        result.insert(
107                            field.name.to_string(),
108                            self.arbitrary_leaf_field(field.ty().inner_named_type())?,
109                        );
110                    } else if field.selection_set.is_empty() && field.ty().is_list() {
111                        result.insert(
112                            field.name.to_string(),
113                            self.repeated_arbitrary_leaf_field(field.ty().inner_named_type())?,
114                        );
115                    } else if !field.selection_set.is_empty() && !field.ty().is_list() {
116                        result.insert(
117                            field.name.to_string(),
118                            self.arbitrary_selection_set(&field.selection_set)?,
119                        );
120                    } else {
121                        result.insert(
122                            field.name.to_string(),
123                            self.repeated_arbitrary_selection_set(&field.selection_set)?,
124                        );
125                    }
126                }
127                Selection::FragmentSpread(fragment) => {
128                    if let Some(fragment_def) = self.doc.fragments.get(&fragment.fragment_name) {
129                        let value = self.arbitrary_selection_set(&fragment_def.selection_set)?;
130                        if let Some(value_obj) = value.as_object() {
131                            result.extend(value_obj.clone());
132                        }
133                    }
134                }
135                Selection::InlineFragment(inline_fragment) => {
136                    let value = self.arbitrary_selection_set(&inline_fragment.selection_set)?;
137                    if let Some(value_obj) = value.as_object() {
138                        result.extend(value_obj.clone());
139                    }
140                }
141            }
142        }
143
144        Ok(Value::Object(result))
145    }
146
147    fn repeated_arbitrary_selection_set(&mut self, selection_set: &SelectionSet) -> Result<Value> {
148        let num_values = self.arbitrary_len()?;
149        let mut values = Vec::with_capacity(num_values);
150        for _ in 0..num_values {
151            values.push(self.arbitrary_selection_set(selection_set)?);
152        }
153        Ok(Value::Array(values))
154    }
155
156    fn arbitrary_leaf_field(&mut self, type_name: &Name) -> Result<Value> {
157        let extended_ty = self.schema.types.get(type_name).unwrap();
158        match extended_ty {
159            ExtendedType::Enum(enum_ty) => {
160                let enum_value = self.u.choose_iter(enum_ty.values.values())?;
161                Ok(Value::String(enum_value.value.to_string().into()))
162            }
163            ExtendedType::Scalar(scalar) => {
164                if scalar.name == "Boolean" {
165                    let random_bool = self.u.arbitrary::<bool>()?;
166                    Ok(Value::Bool(random_bool))
167                } else if scalar.name == "Int" || scalar.name == "ID" {
168                    let random_int = self.u.int_in_range(0..=100)?;
169                    Ok(Value::Number(random_int.into()))
170                } else if scalar.name == "Float" {
171                    let random_float = self.u.arbitrary::<f64>()?;
172                    let Some(random_number) = Number::from_f64(random_float) else {
173                        return Err(arbitrary::Error::IncorrectFormat);
174                    };
175                    Ok(Value::Number(random_number))
176                } else if scalar.name == "String" {
177                    let random_string = self.u.arbitrary::<String>()?;
178                    Ok(Value::String(random_string.into()))
179                } else if let Some(custom_generator) =
180                    self.custom_scalar_generators.get(&scalar.name)
181                {
182                    let random_value = custom_generator(self.u)?;
183                    Ok(random_value)
184                } else {
185                    // Likely a custom scalar which hasn't had a generator registered
186                    let random_string = self.u.arbitrary::<String>()?;
187                    Ok(Value::String(random_string.into()))
188                }
189            }
190            _ => unreachable!(
191                "We are in a field with an empty selection set, so it must be a scalar or enum type"
192            ),
193        }
194    }
195
196    fn repeated_arbitrary_leaf_field(&mut self, type_name: &Name) -> Result<Value> {
197        let num_values = self.arbitrary_len()?;
198        let mut values = Vec::with_capacity(num_values);
199        for _ in 0..num_values {
200            values.push(self.arbitrary_leaf_field(type_name)?);
201        }
202        Ok(Value::Array(values))
203    }
204
205    fn arbitrary_len(&mut self) -> Result<usize> {
206        // Ideally, we would use `u.arbitrary_len()` to ensure we can generate enough values from
207        // the remaining bytes, but it needs a type `T: Arbitrary` which `Value` does not implement.
208        self.u.int_in_range(self.min_list_size..=self.max_list_size)
209    }
210
211    fn should_be_null(&mut self) -> Result<bool> {
212        if let Some((numerator, denominator)) = self.null_ratio {
213            self.u.ratio(numerator, denominator)
214        } else {
215            Ok(false)
216        }
217    }
218}
apollo_smith/response.rs

apollo_smith/
response.rs
