Update generator to consider 201 as success response code

fuzz_lightyear/generator.py

@@ -58,7 +58,7 @@ def generate_sequences(
     #       (rather than starting operation), so that it's clearer for
     #       output.
     client = get_abstraction().client
-    request_graph = _generate_request_graph()
+    request_graph = generate_request_graph()
     for tag_group in get_fuzzable_tags(client):
         last_results = []   # type: List
         for _ in range(n):
@@ -124,7 +124,7 @@ def _add_request_to_sequence(
 
 
 @lru_cache(maxsize=1)
-def _generate_request_graph() -> Dict[str, set]:
+def generate_request_graph() -> Dict[str, set]:
     """
     Generates a directed graph, represented as an adjacency list.
 
@@ -140,6 +140,8 @@ def _generate_request_graph() -> Dict[str, set]:
         for operation_id in dir(getattr(client, tag_group)):
             operation = getattr(getattr(client, tag_group), operation_id).operation
             responses = operation.op_spec.get('responses', {}).get('200', {})
+            if not responses:
+                responses = operation.op_spec.get('responses', {}).get('201', {})
             if responses:
                 response_params = list(
                     responses.get('schema', {}).get('properties', {}).keys(),

testing/vulnerable_app/views/sequence.py

@@ -64,7 +64,7 @@ def get(self):
 @ns.route('/side-effect/create')
 class CreateWithSideEffect(Resource):
     @api.doc(security='apikey')
-    @api.response(200, 'Success', model=widget_model)
+    @api.response(201, 'Success', model=widget_model)
     @requires_user
     def post(self, user):
 
@@ -81,7 +81,7 @@ def post(self, user):
 
         return {
             'id': widget_id,
-        }
+        }, 201
 
 
 @ns.route('/side-effect/get/<int:id>')

tests/integration/generator_test.py

@@ -6,6 +6,7 @@
 import pytest
 
 import fuzz_lightyear
+from fuzz_lightyear.generator import generate_request_graph
 from fuzz_lightyear.generator import generate_sequences
 
 
@@ -166,6 +167,17 @@ def test_length_three(mock_client):
     )
 
 
+def test_request_graph(mock_client):
+    requested_graph = generate_request_graph()
+    node = 'get_bravo_two'
+    expected_edges = [
+        'get_get_with_side_effect_unsafe', 'post_create_with_side_effect',
+        'get_get_with_side_effect_safe',
+    ]
+    for e in expected_edges:
+        assert e in requested_graph[node]
+
+
 def is_in_result(expected_sequence, result):
     for sequence in result:
         formatted_sequence = []