Merge pull request #51 from chisholm/halt-generator-fix

Fix reference graph generator's "halt" generator to inherit customizations

stix2generator/generation/object_generator.py

@@ -1,3 +1,4 @@
+import copy
 import itertools
 import logging
 import math
@@ -106,8 +107,8 @@ def __init__(
 
         :param spec_registry: A name->specification mapping used to look
             up references inside of specifications
-        :param semantic_providers: A list of semantic providers (e.g.
-            instances of subclasses of SemanticsProvider)
+        :param semantic_providers: An iterable of semantic providers (e.g.
+            instances of subclasses of SemanticsProvider), or None
         :param config: A Config instance giving user settings regarding
             generation.  If None, defaults will be used.
         """
@@ -142,6 +143,27 @@ def spec_names(self):
         """
         return self.__specs.keys()
 
+    def derive_generator(self, new_config):
+        """
+        Create an object generator using this generator's registry and
+        semantics providers, but with a given config.
+
+        :param new_config: An instance of Config giving desired configuration
+            settings for the new object generator
+        :return: A new object generator
+        """
+
+        our_providers = copy.deepcopy(set(self.__semantics.values()))
+        our_specs = copy.deepcopy(self.__specs)
+
+        new_generator = ObjectGenerator(
+            spec_registry=our_specs,
+            semantic_providers=our_providers,
+            config=new_config
+        )
+
+        return new_generator
+
     def generate(
         self, spec_name, expected_type=None, spec_name_stack=None,
         value_constraint=None

stix2generator/generation/reference_graph_generator.py

@@ -515,7 +515,7 @@ def __init__(
             .Config(
                 **halt_generator_config_dict
             )
-        self.__halt_generator = stix2generator.create_object_generator(
+        self.__halt_generator = object_generator.derive_generator(
             halt_generator_config
         )
 

stix2generator/test/test_reference_graph_generator.py

@@ -1,12 +1,14 @@
 import copy
 
 import pytest
+import stix2
 import stix2.base
 import stix2.utils
 
 import stix2generator
 import stix2generator.generation.object_generator
 import stix2generator.generation.reference_graph_generator
+import stix2generator.generation.semantics
 import stix2generator.test.utils
 
 _TLP_MARKING_DEFINITION_IDS = {
@@ -959,3 +961,88 @@ def test_not_parsing(num_trials):
 
             else:
                 assert isinstance(obj, dict)
+
+
+def test_ref_gen_with_custom():
+    """
+    Set up a reference graph generator with an object generator with stuff
+    you would not get by default via
+    stix2generator.create_object_generator(...).  Reference graph generator
+    creates a derivative "halt" generator, and we need to ensure that the
+    latter inherits the same config as the original object generator, including
+    registry and semantics providers.
+    """
+    class TestSemantics(stix2generator.generation.semantics.SemanticsProvider):
+        def get_semantics(self):
+            return ["testxyz"]
+
+        def testxyz(self, spec, generator, constraint):
+            return "test"
+
+    custom_registry = {
+        # A self-referential type, to cause simple reference chains
+        "test-ref-type": {
+            "type": "object",
+            "optional": ["obj_ref"],
+            "properties": {
+                "id": {
+                    "type": "string",
+                    "semantics": "stix-id",
+                    "stix-type": "test-ref-type"
+                },
+                "type": "test-ref-type",
+                "test_prop": {
+                    "type": "string",
+                    "semantics": "testxyz"
+                },
+                "obj_ref": {
+                    "type": "string",
+                    "semantics": "stix-id",
+                    "stix-type": "test-ref-type"
+                }
+            }
+        }
+    }
+
+    @stix2.CustomObject("test-ref-type", [])
+    class TestRefType:
+        # This class won't be used since we'll turn the parse setting off; but
+        # we need a registration so test-ref-type is seen as a generatable
+        # type.  So we can leave it empty.
+        pass
+
+    semantics_providers = [
+        TestSemantics(),
+        stix2generator.generation.semantics.STIXSemantics()
+    ]
+
+    obj_gen_config = stix2generator.generation.object_generator.Config(
+        optional_property_probability=1,
+        minimize_ref_properties=False
+    )
+
+    obj_gen = stix2generator.generation.object_generator.ObjectGenerator(
+        custom_registry, semantics_providers, obj_gen_config
+    )
+
+    ref_gen_config = stix2generator.generation.reference_graph_generator \
+        .Config(
+            graph_type="TREE",
+            parse=False
+        )
+
+    ref_gen = stix2generator.generation.reference_graph_generator \
+        .ReferenceGraphGenerator(
+            obj_gen, ref_gen_config
+        )
+
+    _, graph = ref_gen.generate()
+
+    for obj in graph.values():
+        # Ensure our semantics provider was invoked properly for all objects
+        assert obj["test_prop"] == "test"
+
+        # Ensure no dangling references
+        obj_ref = obj.get("obj_ref")
+        if obj_ref is not None:
+            assert obj_ref in graph