Refactor and test init_additional_STUN_servers.

* Make it accept any server sequence. * Make it idempotent.
Vorstellung · Oct 28, 2014 · ecd53dc · ecd53dc
1 parent 133ee9d
commit ecd53dc
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 22 deletions.
diff --git a/node/network_util.py b/node/network_util.py
@@ -4,28 +4,32 @@
 import stun
 
 
-def init_additional_STUN_servers():
-    """try calling this method only once"""
-    # list of additional stun servers taken (and tested) from natvpn project
-    # https://code.google.com/p/natvpn/source/browse/trunk/stun_server_list
-    # removed those that didn't ping back.
-    stun.stun_servers_list = stun.stun_servers_list + (
-        'stun.l.google.com',
-        'stun1.l.google.com',
-        'stun2.l.google.com',
-        'stun3.l.google.com',
-        'stun4.l.google.com',
-        'stun.ekiga.net',
-        'stun.ideasip.com',
-        'stun.iptel.org',
-        'stun.schlund.de',
-        'stunserver.org',
-        'stun.voiparound.com',
-        'stun.voipbuster.com',
-        'stun.voipstunt.com',
-        'stun.voxgratia.org',
-        'stun.xten.com'
-    )
+# List taken and tested from natvpn project:
+# https://code.google.com/p/natvpn/source/browse/trunk/stun_server_list
+_ADDITIONAL_STUN_SERVERS = (
+    'stun.l.google.com',
+    'stun1.l.google.com',
+    'stun2.l.google.com',
+    'stun3.l.google.com',
+    'stun4.l.google.com',
+    'stun.ekiga.net',
+    'stun.ideasip.com',
+    'stun.iptel.org',
+    'stun.schlund.de',
+    'stunserver.org',
+    'stun.voiparound.com',
+    'stun.voipbuster.com',
+    'stun.voipstunt.com',
+    'stun.voxgratia.org',
+    'stun.xten.com'
+)
+
+
+def init_additional_STUN_servers(servers=_ADDITIONAL_STUN_SERVERS):
+    """Inject list of additional STUN servers."""
+    server_set = set(stun.stun_servers_list)
+    server_set.update(servers)
+    stun.stun_servers_list = tuple(server_set)
 
 
 def check_NAT_status():

diff --git a/test/test_network_util.py b/test/test_network_util.py
@@ -1,9 +1,33 @@
+import collections
 import unittest
 
+import stun
+
 from node import network_util
 
 
 class TestNodeNetworkUtil(unittest.TestCase):
+
+    def test_init_additional_stun_servers(self):
+        stun_servers_pre = stun.stun_servers_list
+        new_stun_servers = (
+            'stun.openbazaar1.com',
+            'stun.openbazaar2.com'
+        )
+        network_util.init_additional_STUN_servers(servers=new_stun_servers)
+
+        counter = collections.Counter(stun.stun_servers_list)
+
+        # Check all new STUN servers are in.
+        for server in new_stun_servers:
+            self.assertEqual(counter[server], 1)
+
+        network_util.init_additional_STUN_servers(servers=new_stun_servers)
+
+        # Check no STUN server was removed or added twice.
+        for server in stun_servers_pre:
+            self.assertEqual(counter[server], 1)
+
     def test_is_loopback_addr(self):
         self.assertTrue(network_util.is_loopback_addr("127.0.0.1"))
         self.assertTrue(network_util.is_loopback_addr("localhost"))