From 4d6c82b4cd06453f35f538b8a19c13a3a7afb204 Mon Sep 17 00:00:00 2001
From: Brandon Rodriguez <brodriguez8774@gmail.com>
Date: Thu, 5 Dec 2019 04:09:13 -0500
Subject: [PATCH] Create method to find "largest connected capacity" neighbor
 node

---
 .../graphs/network_flow_graph/components.py   | 35 ++++++++
 .../graphs/network_flow_graph/node.py         | 86 +++++++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/resources/graphs/network_flow_graph/components.py b/resources/graphs/network_flow_graph/components.py
index deda0f7..0a37bc4 100644
--- a/resources/graphs/network_flow_graph/components.py
+++ b/resources/graphs/network_flow_graph/components.py
@@ -105,3 +105,38 @@ class NetworkFlowNode(DirectedWeightedNode):
         # This is necessary for inheritance, or else child classes will only have access to parent functions.
         self._edge_type = NetworkFlowEdge
         self._node_type = NetworkFlowNode
+
+    def get_highest_capacity(self, incoming_only=False, outgoing_only=False):
+        """
+        Finds node connection with highest capacity out of all neighbors.
+        :param incoming_only: Bool indicating if only incoming edges should be examined.
+        :param outgoing_only: Bool indicating if only outgoing edges should be examined.
+        :return: Node with highest capacity connection | None if no connections.
+        """
+        # Validate args.
+        if incoming_only and outgoing_only:
+            raise ValueError('Only one of "incoming_only" or "outgoing_only" should be True. Found both True.')
+
+        capacity_node = None
+        capacity_edge = None
+
+        # Get edge list based on args.
+        if incoming_only:
+            # Only check incoming edges.
+            edge_list = self.get_incoming_edges()
+        elif outgoing_only:
+            # Only check outgoing edges.
+            edge_list = self.get_outgoing_edges()
+        else:
+            # Get all edges.
+            edge_list = self.get_edges().values()
+
+        # Loop through all edge connections.
+        for edge_connection in edge_list:
+            # Check if connection has higher capacity than current best.
+            if capacity_node is None or edge_connection.get_capacity() > capacity_edge.get_capacity():
+                capacity_edge = edge_connection
+                capacity_node = edge_connection.get_partner_node(self)
+
+        # Return highest found node.
+        return capacity_node
diff --git a/tests/resources/graphs/network_flow_graph/node.py b/tests/resources/graphs/network_flow_graph/node.py
index 4c743d6..0353130 100644
--- a/tests/resources/graphs/network_flow_graph/node.py
+++ b/tests/resources/graphs/network_flow_graph/node.py
@@ -22,3 +22,89 @@ class TestNetworkFlowNode(unittest.TestCase):
     def test__node_initialization(self):
         self.assertEqual(self.test_node._edge_type, NetworkFlowEdge)
         self.assertEqual(self.test_node._node_type, NetworkFlowNode)
+
+    def test__get_highest_capacity(self):
+        # Create additional nodes.
+        node_1 = NetworkFlowNode(1)
+        node_2 = NetworkFlowNode(2)
+        node_3 = NetworkFlowNode(3)
+        node_4 = NetworkFlowNode(4)
+
+        with self.subTest('With all edge directions.'):
+            # Test with no connections.
+            self.assertIsNone(self.test_node.get_highest_capacity())
+
+            # Test with one connection.
+            edge_1 = self.test_node.connect_node(node_1)
+            edge_1.set_capacity(15)
+            self.assertEqual(self.test_node.get_highest_capacity(), node_1)
+
+            # Test with two connections.
+            edge_2 = self.test_node.connect_node(node_2)
+            edge_2.set_capacity(12)
+            self.assertEqual(self.test_node.get_highest_capacity(), node_1)
+
+            # Test with three connections.
+            edge_3 = self.test_node.connect_node(node_3)
+            edge_3.set_capacity(13)
+            self.assertEqual(self.test_node.get_highest_capacity(), node_1)
+
+            # Test with a new higher capacity, to ensure it wasn't just taking the first valid node.
+            edge_4 = self.test_node.connect_node(node_4)
+            edge_4.set_capacity(20)
+            self.assertEqual(self.test_node.get_highest_capacity(), node_4)
+
+            # Reset edge connections for following subtests.
+            self.test_node.disconnect_node(node_identifier=node_1)
+            self.test_node.disconnect_node(node_identifier=node_2)
+            self.test_node.disconnect_node(node_identifier=node_3)
+            self.test_node.disconnect_node(node_identifier=node_4)
+
+        with self.subTest('With only incoming edges.'):
+            # Test with no connections.
+            self.assertIsNone(self.test_node.get_highest_capacity(incoming_only=True))
+
+            # Test with incoming connection.
+            edge_1 = node_1.connect_node(self.test_node)
+            edge_1.set_capacity(10)
+            self.assertEqual(self.test_node.get_highest_capacity(incoming_only=True), node_1)
+
+            # Test with two incoming connections.
+            edge_2 = node_2.connect_node(self.test_node)
+            edge_2.set_capacity(12)
+            self.assertEqual(self.test_node.get_highest_capacity(incoming_only=True), node_2)
+
+            # Test with outgoing connection.
+            edge_3 = self.test_node.connect_node(node_3)
+            edge_3.set_capacity(15)
+            self.assertEqual(self.test_node.get_highest_capacity(incoming_only=True), node_2)
+
+            # Reset edge connections for following subtests.
+            self.test_node.disconnect_node(node_identifier=node_1)
+            self.test_node.disconnect_node(node_identifier=node_2)
+            self.test_node.disconnect_node(node_identifier=node_3)
+
+        with self.subTest('With only outgoing edges.'):
+            # Test with no connections.
+            self.assertIsNone(self.test_node.get_highest_capacity(outgoing_only=True))
+
+            # Test with outgoing connection.
+            edge_1 = self.test_node.connect_node(node_1)
+            edge_1.set_capacity(10)
+            self.assertEqual(self.test_node.get_highest_capacity(outgoing_only=True), node_1)
+
+            # Test with two outgoing connections.
+            edge_2 = self.test_node.connect_node(node_2)
+            edge_2.set_capacity(12)
+            self.assertEqual(self.test_node.get_highest_capacity(outgoing_only=True), node_2)
+
+            # Test with incoming connection.
+            edge_3 = node_3.connect_node(self.test_node)
+            edge_3.set_capacity(15)
+            self.assertEqual(self.test_node.get_highest_capacity(outgoing_only=True), node_2)
+
+            # Reset edge connections for following subtests.
+            self.test_node.disconnect_node(node_identifier=node_1)
+            self.test_node.disconnect_node(node_identifier=node_2)
+            self.test_node.disconnect_node(node_identifier=node_3)
+
-- 
GitLab