Haircontrol: Small refactoring

* haircontrol/discovery.py (Discovery.compute_neighbourhood.set_uplink):
  Test if link has an uplink first, so as to avoid code duplication.

1 files changed, 8 insertions, 10 deletions

diff --git a/haircontrol/discovery.py b/haircontrol/discovery.py
index 70bdc6a..991abc3 100644
--- a/haircontrol/discovery.py
+++ b/haircontrol/discovery.py
@@ -242,10 +242,11 @@ class Discovery:
             return True
 
         def set_uplink (ip, interface):
-            # Set IP as a direct neighbour of INTERFACE. Also signal
-            # that equipment relative to IP has a known uplink.
-            interface.direct_neighbours.append(ip)
-            with_uplink.append(ip)
+            # Set IP as a direct neighbour of INTERFACE, if it doesn't
+            # have an uplink already.
+            if ip not in with_uplink:
+                interface.direct_neighbours.append(ip)
+                with_uplink.append(ip)
 
         def link_to_outer_net (equipment):
             # Link EQUIPMENT with outer network. Orphans in the
@@ -253,10 +254,8 @@ class Discovery:
             for interface in equipment.ifaces.values():
                 if is_outer_interface(interface):
                     for mac in interface.mac_seen:
-                        if mac not in mac_to_ip: continue
-                        ip = mac_to_ip[mac]
-                        if ip not in with_uplink:
-                            set_uplink(ip, interface)
+                        if mac in mac_to_ip:
+                            set_uplink(mac_to_ip[mac], interface)
 
         ### Main
 
@@ -290,5 +289,4 @@ class Discovery:
             outer_net.extend(current_level)
             for ip in current_level: del equipments[ip]
         exit_iface = self.net.equipments[gateway_ip].ifaces[gateway_iface_name]
-        for ip in outer_net:
-            if ip not in with_uplink: set_uplink(ip, exit_iface)
+        for ip in outer_net: set_uplink(ip, exit_iface)