From 2452cd414f692ccc28aa3ec4827a71eca4b28d96 Mon Sep 17 00:00:00 2001
From: Andrew Dickinson <andrewjd@amazon.com>
Date: Fri, 17 Jan 2025 21:28:09 -0500
Subject: [PATCH] Pick up existing node on install create

---
 src/meshapi/models/install.py     |  7 +++
 src/meshapi/tests/test_install.py | 74 ++++++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/src/meshapi/models/install.py b/src/meshapi/models/install.py
index e50d6b62..b6f265e2 100644
--- a/src/meshapi/models/install.py
+++ b/src/meshapi/models/install.py
@@ -142,4 +142,11 @@ def save(self, *args: Any, **kwargs: Any) -> None:
             if self.install_number != original.install_number:
                 raise ValidationError("Install number is immutable")
 
+        if self._state.adding:
+            if not self.node:
+                if self.building.primary_node:
+                    self.node = self.building.primary_node
+                elif self.building.nodes:
+                    self.node = self.building.nodes.first()
+
         super().save(*args, **kwargs)
diff --git a/src/meshapi/tests/test_install.py b/src/meshapi/tests/test_install.py
index b79494f5..9178fefe 100644
--- a/src/meshapi/tests/test_install.py
+++ b/src/meshapi/tests/test_install.py
@@ -6,7 +6,7 @@
 from django.core.exceptions import ValidationError
 from django.test import TestCase
 
-from meshapi.models import Building, Install, Member
+from meshapi.models import Building, Install, Member, Node
 from meshapi.tests.sample_data import sample_building, sample_install, sample_member
 
 
@@ -43,6 +43,78 @@ def test_construct_install_no_id_yes_install_number(self):
         self.assertIsNotNone(install2.id)
         self.assertEqual(install2.install_number, 89)
 
+    def test_construct_install_building_has_node(self):
+        building_node = Node(status=Node.NodeStatus.ACTIVE, latitude=0, longitude=0)
+        building_node.save()
+
+        self.building_1.primary_node = building_node
+        self.building_1.save()
+
+        install = Install(**self.sample_install_copy, install_number=45)
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, building_node)
+
+    def test_construct_install_building_has_node_without_primary(self):
+        building_node = Node(status=Node.NodeStatus.ACTIVE, latitude=0, longitude=0)
+        building_node.save()
+
+        self.building_1.nodes.add(building_node)
+        self.building_1.save()
+
+        install = Install(**self.sample_install_copy, install_number=45)
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, building_node)
+
+    def test_remove_node_when_building_has_one(self):
+        building_node = Node(status=Node.NodeStatus.ACTIVE, latitude=0, longitude=0)
+        building_node.save()
+
+        self.building_1.primary_node = building_node
+        self.building_1.save()
+
+        install = Install(**self.sample_install_copy, install_number=45)
+        install.node = building_node
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, building_node)
+
+        install.node = None
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, None)
+
+    def test_set_install_to_different_to_building(self):
+        building1_node = Node(status=Node.NodeStatus.ACTIVE, latitude=0, longitude=0)
+        building2_node = Node(status=Node.NodeStatus.ACTIVE, latitude=0, longitude=0)
+        building1_node.save()
+        building2_node.save()
+
+        self.building_1.primary_node = building1_node
+        self.building_1.save()
+
+        building2 = Building(**sample_building)
+        building2.primary_node = building2_node
+        building2.save()
+
+        install = Install(**self.sample_install_copy, install_number=45)
+        install.node = building1_node
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, building1_node)
+
+        install.building = building2
+        install.save()
+
+        install.refresh_from_db()
+        self.assertEqual(install.node, building1_node)
+
     def test_construct_install_yes_id_no_install_number(self):
         install = Install(
             id=uuid.UUID("23ef170c-f37d-44e3-aaac-93dae636c86e"),