Browse Source

fixed get_nth_element and added tests

get_nth_element can now handle one list, empty lists and zero lists without crashing
added testcases for these cases
Stefano Acquaviti 6 years ago
parent
commit
bb77cb89b7
2 changed files with 23 additions and 9 deletions
  1. 13 5
      code/ID2TLib/Utility.py
  2. 10 4
      code/Test/test_Utility.py

+ 13 - 5
code/ID2TLib/Utility.py

@@ -1,7 +1,7 @@
 import ipaddress
 import os
 
-from random import randint, uniform
+from random import randint, uniform, choice
 from os import urandom
 from datetime import datetime
 from calendar import timegm
@@ -71,10 +71,18 @@ def get_nth_random_element(*element_list):
     :param element_list: An arbitrary number of lists.
     :return: A tuple of the n-th element of every list.
     """
-    range_max = min([len(x) for x in element_list])
-    if range_max > 0: range_max -= 1
-    n = randint(0, range_max)
-    return tuple(x[n] for x in element_list)
+    if len(element_list) <= 0:
+        return None
+    elif len(element_list) == 1 and len(element_list[0]) > 0:
+        return choice(element_list[0])
+    else:
+        range_max = min([len(x) for x in element_list])
+        if range_max > 0:
+            range_max -= 1
+            n = randint(0, range_max)
+            return tuple(x[n] for x in element_list)
+        else:
+            return None
 
 
 def index_increment(number: int, max: int):

+ 10 - 4
code/Test/test_Utility.py

@@ -40,10 +40,16 @@ class TestUtility(unittest.TestCase):
         results = [("A", 1), ("B", 2)]
         self.assertIn(Utility.get_nth_random_element(letters, numbers), results)
 
-    # TODO: ???
-    #def test_get_nth_random_element_single_list(self):
-        #letters = ["A", "B", "C"]
-        #self.assertIn(Utility.get_nth_random_element(letters), letters)
+    def test_get_nth_random_element_single_list(self):
+        letters = ["A", "B", "C"]
+        self.assertIn(Utility.get_nth_random_element(letters), letters)
+
+    def test_get_nth_random_element_empty_list(self):
+        letters = ["A", "B", "C"]
+        self.assertEqual(Utility.get_nth_random_element(letters, []), None)
+
+    def test_get_nth_random_element_nothing(self):
+        self.assertEqual(Utility.get_nth_random_element(), None)
 
     def test_index_increment_not_max(self):
         self.assertEqual(Utility.index_increment(5, 10), 6)