Tests

nautobot_design_builder/contrib/tests/test_ext.py

@@ -2,11 +2,10 @@
 
 import os
 
-from django.test import TestCase
+from nautobot_design_builder.tests.test_builder import BuilderTestCase
 
-from nautobot_design_builder.tests.test_builder import builder_test_case
 
-
-@builder_test_case(os.path.join(os.path.dirname(__file__), "testdata"))
-class TestAgnosticExtensions(TestCase):
+class TestAgnosticExtensions(BuilderTestCase):
     """Test contrib extensions against any version of Nautobot."""
+
+    data_dir = os.path.join(os.path.dirname(__file__), "testdata")

nautobot_design_builder/fields.py

@@ -213,25 +213,49 @@ class ManyToManyField(BaseModelField, RelationshipFieldMixin):  # pylint:disable
 
     def __init__(self, field: django_models.Field):  # noqa:D102
         super().__init__(field)
-        if hasattr(field.remote_field, "through"):
-            through = field.remote_field.through
-            if not through._meta.auto_created:
-                self.related_model = through
+        self.auto_through = True
+        self.through_fields = field.remote_field.through_fields
+        through = field.remote_field.through
+        if not through._meta.auto_created:
+            self.auto_through = False
+            self.related_model = through
+            if field.remote_field.through_fields:
+                self.link_field = field.remote_field.through_fields[0]
+            else:
+                for f in through._meta.fields:
+                    if f.related_model == field.model:
+                        self.link_field = f.name
 
     @debug_set
     def __set__(self, obj: "ModelInstance", values):  # noqa:D105
         def setter():
             items = []
             for value in values:
                 value = self._get_instance(obj, value, getattr(obj.instance, self.field_name))
+                if self.auto_through:
+                    # Only need to call `add` if the through relationship was
+                    # auto-created. Otherwise we explicitly create the through
+                    # object
+                    items.append(value.instance)
+                else:
+                    setattr(value.instance, self.link_field, obj.instance)
                 if value.metadata.created:
                     value.save()
-                items.append(value.instance)
-            getattr(obj.instance, self.field_name).add(*items)
+                else:
+                    value.environment.journal.log(value)
+            if items:
+                getattr(obj.instance, self.field_name).add(*items)
 
         obj.connect("POST_INSTANCE_SAVE", setter)
 
 
+class ManyToManyRelField(ManyToManyField):  # pylint:disable=too-few-public-methods
+    """Reverse many to many relationship field."""
+
+    def __init__(self, field: django_models.Field):  # noqa:D102
+        super().__init__(field.remote_field)
+
+
 class GenericRelationField(BaseModelField, RelationshipFieldMixin):  # pylint:disable=too-few-public-methods
     """Generic relationship field."""
 
@@ -259,13 +283,30 @@ def __set__(self, obj: "ModelInstance", value):  # noqa:D105
         setattr(obj.instance, ct_field, ContentType.objects.get_for_model(value.instance))
 
 
-class TagField(ManyToManyField):  # pylint:disable=too-few-public-methods
+class TagField(BaseModelField, RelationshipFieldMixin):  # pylint:disable=too-few-public-methods
     """Taggit field."""
 
     def __init__(self, field: django_models.Field):  # noqa:D102
         super().__init__(field)
         self.related_model = field.remote_field.model
 
+    def __set__(self, obj: "ModelInstance", values):  # noqa:D105
+        # I hate that this code is almost identical to the ManyToManyField
+        # __set__ code, but I don't see an easy way to DRY it up at the
+        # moment.
+        def setter():
+            items = []
+            for value in values:
+                value = self._get_instance(obj, value, getattr(obj.instance, self.field_name))
+                if value.metadata.created:
+                    value.save()
+                else:
+                    value.environment.journal.log(value)
+                items.append(value.instance)
+            getattr(obj.instance, self.field_name).add(*items)
+
+        obj.connect("POST_INSTANCE_SAVE", setter)
+
 
 class GenericRelField(BaseModelField, RelationshipFieldMixin):  # pylint:disable=too-few-public-methods
     """Field used as part of content-types generic relation."""
@@ -348,8 +389,10 @@ def field_factory(arg1, arg2) -> ModelField:
         field = ForeignKeyField(arg2)
     elif isinstance(arg2, django_models.ManyToOneRel):
         field = ManyToOneRelField(arg2)
-    elif isinstance(arg2, (django_models.ManyToManyField, django_models.ManyToManyRel)):
+    elif isinstance(arg2, django_models.ManyToManyField):
         field = ManyToManyField(arg2)
+    elif isinstance(arg2, django_models.ManyToManyRel):
+        field = ManyToManyRelField(arg2)
     else:
         raise DesignImplementationError(f"Cannot manufacture field for {type(arg2)}, {arg2} {arg2.is_relation}")
     return field

nautobot_design_builder/management/commands/install_demo_designs.py

@@ -8,13 +8,22 @@
 class Command(BaseCommand):
     """Create a git datasource pointed to the demo designs repo."""
 
+    def add_arguments(self, parser):
+        """Add the branch argument to the command."""
+        parser.add_argument(
+            "--branch",
+            action="store",
+            help="Specify which branch to use in the demo-design repository (default: main).",
+            default="main",
+        )
+
     def handle(self, *args, **options):
         """Handle the execution of the command."""
         GitRepository.objects.get_or_create(
             name="Demo Designs",
             defaults={
                 "remote_url": "https://github.com/nautobot/demo-designs.git",
-                "branch": "main",
+                "branch": options["branch"],
                 "provided_contents": ["extras.job"],
             },
         )

nautobot_design_builder/tests/test_builder.py

@@ -3,7 +3,7 @@
 import importlib
 from operator import attrgetter
 import os
-from unittest.mock import Mock, patch
+from unittest.mock import patch
 import yaml
 
 from django.db.models import Manager, Q
@@ -64,6 +64,24 @@ def check_model_not_exist(test, check, index):
         values = _get_value(check)
         test.assertEqual(len(values), 0, msg=f"Check {index}")
 
+    @staticmethod
+    def check_in(test, check, index):
+        """Check that a model does not exist."""
+        value0 = _get_value(check[0])[0]
+        value1 = _get_value(check[1])
+        if len(value1) == 1:
+            value1 = value1[0]
+        test.assertIn(value0, value1, msg=f"Check {index}")
+
+    @staticmethod
+    def check_not_in(test, check, index):
+        """Check that a model does not exist."""
+        value0 = _get_value(check[0])[0]
+        value1 = _get_value(check[1])
+        if len(value1) == 1:
+            value1 = value1[0]
+        test.assertNotIn(value0, value1, msg=f"Check {index}")
+
 
 def _get_value(check_info):
     if "value" in check_info:
@@ -105,48 +123,69 @@ def _testcases(data_dir):
                 yield yaml.safe_load(file), filename
 
 
-def builder_test_case(data_dir):
-    """Decorator to load tests into a TestCase from a data directory."""
+class _BuilderTestCaseMeta(type):
+    def __new__(mcs, name, bases, dct):
+        cls = super().__new__(mcs, name, bases, dct)
+        data_dir = getattr(cls, "data_dir", None)
+        if data_dir is None:
+            return cls
 
-    def class_wrapper(test_class):
         for testcase, filename in _testcases(data_dir):
+            if testcase.get("abstract", False):
+                continue
             # Strip the .yaml extension
             testcase_name = f"test_{filename[:-5]}"
 
             # Create a new closure for testcase
             def test_wrapper(testcase):
-                @patch("nautobot_design_builder.design.Environment.roll_back")
-                def test_runner(self, roll_back: Mock):
+                def test_runner(self: "BuilderTestCase"):
                     if testcase.get("skip", False):
                         self.skipTest("Skipping due to testcase skip=true")
-                    extensions = []
-                    for extension in testcase.get("extensions", []):
-                        extensions.append(_load_class(extension))
-
-                    with self.captureOnCommitCallbacks(execute=True):
-                        for design in testcase["designs"]:
-                            environment = Environment(extensions=extensions)
-                            commit = design.pop("commit", True)
-                            environment.implement_design(design=design, commit=commit)
-                            if not commit:
-                                roll_back.assert_called()
-
-                    for index, check in enumerate(testcase.get("checks", [])):
-                        for check_name, args in check.items():
-                            _check_name = f"check_{check_name}"
-                            if hasattr(BuilderChecks, _check_name):
-                                getattr(BuilderChecks, _check_name)(self, args, index)
-                            else:
-                                raise ValueError(f"Unknown check {check_name} {check}")
+                    self._run_test_case(testcase, cls.data_dir)  # pylint:disable=protected-access
 
                 return test_runner
 
-            setattr(test_class, testcase_name, test_wrapper(testcase))
-        return test_class
+            setattr(cls, testcase_name, test_wrapper(testcase))
+        return cls
+
+
+class BuilderTestCase(TestCase, metaclass=_BuilderTestCaseMeta):  # pylint:disable=missing-class-docstring
+    def _run_checks(self, checks):
+        for index, check in enumerate(checks):
+            for check_name, args in check.items():
+                _check_name = f"check_{check_name}"
+                if hasattr(BuilderChecks, _check_name):
+                    getattr(BuilderChecks, _check_name)(self, args, index)
+                else:
+                    raise ValueError(f"Unknown check {check_name} {check}")
+
+    def _run_test_case(self, testcase, data_dir):
+        with patch("nautobot_design_builder.design.Environment.roll_back") as roll_back:
+            self._run_checks(testcase.get("pre_checks", []))
 
-    return class_wrapper
+            depends_on = testcase.pop("depends_on", None)
+            if depends_on:
+                depends_on_path = os.path.join(data_dir, depends_on)
+                depends_on_dir = os.path.dirname(depends_on_path)
+                with open(depends_on_path, encoding="utf-8") as file:
+                    self._run_test_case(yaml.safe_load(file), depends_on_dir)
 
+            extensions = []
+            for extension in testcase.get("extensions", []):
+                extensions.append(_load_class(extension))
 
-@builder_test_case(os.path.join(os.path.dirname(__file__), "testdata"))
-class TestGeneralDesigns(TestCase):
+            with self.captureOnCommitCallbacks(execute=True):
+                for design in testcase["designs"]:
+                    environment = Environment(extensions=extensions)
+                    commit = design.pop("commit", True)
+                    environment.implement_design(design=design, commit=commit)
+                    if not commit:
+                        roll_back.assert_called()
+
+            self._run_checks(testcase.get("checks", []))
+
+
+class TestGeneralDesigns(BuilderTestCase):
     """Designs that should work with all versions of Nautobot."""
+
+    data_dir = os.path.join(os.path.dirname(__file__), "testdata")

nautobot_design_builder/tests/testdata/assign_tags_by_name.yaml

@@ -0,0 +1,19 @@
+---
+depends_on: "base_test.yaml"
+designs:
+  - tags:
+      - name: "Test Tag"
+        description: "Some Description"
+
+    locations:
+      - name: "site_1"
+        location_type__name: "Site"
+        status__name: "Active"
+        tags:
+          - {"!get:name": "Test Tag"}
+checks:
+  - equal:
+      - model: "nautobot.dcim.models.Location"
+        query: {name: "site_1"}
+        attribute: "tags"
+      - model: "nautobot.extras.models.Tag"

nautobot_design_builder/tests/testdata/assign_tags_by_ref.yaml

@@ -0,0 +1,20 @@
+---
+depends_on: "base_test.yaml"
+designs:
+  - tags:
+      - name: "Test Tag"
+        "!ref": "test_tag"
+        description: "Some Description"
+
+    locations:
+      - name: "site_1"
+        location_type__name: "Site"
+        status__name: "Active"
+        tags:
+          - "!ref:test_tag"
+checks:
+  - equal:
+      - model: "nautobot.dcim.models.Location"
+        query: {name: "site_1"}
+        attribute: "tags"
+      - model: "nautobot.extras.models.Tag"

nautobot_design_builder/tests/testdata/base_test.yaml

@@ -0,0 +1,42 @@
+---
+abstract: true
+designs:
+  - manufacturers:
+      - name: "manufacturer1"
+
+    device_types:
+      - manufacturer__name: "manufacturer1"
+        model: "model name"
+        u_height: 1
+
+    roles:
+      - name: "device role"
+        content_types:
+          - "!get:app_label": "dcim"
+            "!get:model": "device"
+
+    location_types:
+      - name: "Site"
+        content_types:
+          - "!get:app_label": "circuits"
+            "!get:model": "circuittermination"
+          - "!get:app_label": "dcim"
+            "!get:model": "device"
+          - "!get:app_label": "dcim"
+            "!get:model": "powerpanel"
+          - "!get:app_label": "dcim"
+            "!get:model": "rack"
+          - "!get:app_label": "dcim"
+            "!get:model": "rackgroup"
+          - "!get:app_label": "ipam"
+            "!get:model": "prefix"
+          - "!get:app_label": "ipam"
+            "!get:model": "vlan"
+          - "!get:app_label": "ipam"
+            "!get:model": "vlangroup"
+          - "!get:app_label": "virtualization"
+            "!get:model": "cluster"
+        locations:
+          - "name": "Site"
+            "location_type__name": "Site"
+            "status__name": "Active"