test: locker problem

Author: Justintime50

algorithms/sequences/lockers.py

@@ -20,15 +20,27 @@
 
 STUDENTS = 100
 
-lockers = [0] * 100
-
-for student in range(1, STUDENTS + 1):
-    for i, locker in enumerate(lockers):
-        if (i + 1) % student == 0:
-            if lockers[i] == 0:
-                lockers[i] = 1
-            elif lockers[i] == 1:
-                lockers[i] = 0
-
-open_lockers = [i + 1 for i, locker in enumerate(lockers) if locker == 1]
-print(open_lockers)
+
+def solve_locker_problem():
+    lockers = [0] * 100
+
+    for student in range(1, STUDENTS + 1):
+        for i, locker in enumerate(lockers):
+            if (i + 1) % student == 0:
+                if lockers[i] == 0:
+                    lockers[i] = 1
+                elif lockers[i] == 1:
+                    lockers[i] = 0
+
+    open_lockers = [i + 1 for i, locker in enumerate(lockers) if locker == 1]
+
+    return open_lockers
+
+
+def main():
+    solution = solve_locker_problem()
+    print(solution)
+
+
+if __name__ == '__main__':
+    main()

test/unit/sequences/test_lockers.py

@@ -0,0 +1,7 @@
+from algorithms.sequences.lockers import solve_locker_problem
+
+
+def test_solve_locker_problem():
+    solution = solve_locker_problem()
+
+    assert solution == [1, 4, 9, 16, 25, 36, 49, 64, 81, 100]