Updated after tests with real Enigma machine

[cipher-training.git] / test_enigma.py

diff --git a/test_enigma.py b/test_enigma.py
index 830bb2b1fdda74d4a415396edff434c2afc97052..9b59edb392d25fd5daebed18c9964fd34e549f7e 100644 (file)
--- a/test_enigma.py
+++ b/test_enigma.py
@@ -93,6 +93,29 @@ class SimpleWheelTest(unittest.TestCase):
                 for l in string.ascii_lowercase),
             'ajpczwrlfbdkotyuqgenhxmivs')
 
+    def test_set_position(self):
+        wheel_3 = SimpleWheel(wheel_iii_spec)
+        wheel_3.set_position('a')
+        self.assertEqual(wheel_3.position, 0)
+        self.assertEqual(wheel_3.position_l, 'a')
+
+        wheel_3.set_position('f')
+        self.assertEqual(wheel_3.position, 5)
+        self.assertEqual(wheel_3.position_l, 'f')
+
+        wheel_3.advance()
+        self.assertEqual(wheel_3.position, 6)
+        self.assertEqual(wheel_3.position_l, 'g')
+
+        wheel_3.set_position(12)
+        self.assertEqual(wheel_3.position, 12)
+        self.assertEqual(wheel_3.position_l, 'm')
+
+        wheel_3.advance()
+        self.assertEqual(wheel_3.position, 13)
+        self.assertEqual(wheel_3.position_l, 'n')
+
+
     def test_advance(self):
         wheel_3 = SimpleWheel(wheel_iii_spec)
         wheel_3.set_position('a')
@@ -131,7 +154,7 @@ class SimpleWheelTest(unittest.TestCase):
             'tagbpcsdqeufvnzhyixjwlrkom')
 
 
-class SimpleWheelTest(unittest.TestCase):
+class WheelTest(unittest.TestCase):
     def test_init1(self):
         wheel = Wheel(wheel_iii_spec, wheel_iii_pegs, position='b', 
             ring_setting=1)
@@ -158,7 +181,6 @@ class SimpleWheelTest(unittest.TestCase):
                 for l in string.ascii_lowercase),
             'ptlyrmidoxbswhnfckquzgeavj')
 
-
     def test_advance(self):
         wheel = Wheel(wheel_vi_spec, wheel_vi_pegs, position='b', 
             ring_setting=3)
@@ -243,7 +265,6 @@ class SimpleWheelTest(unittest.TestCase):
                 for l in string.ascii_lowercase),
             'ptlyrmidoxbswhnfckquzgeavj')
 
-
     def test_advance_27(self):
         wheel = Wheel(wheel_vi_spec, wheel_vi_pegs, position='b', 
             ring_setting=3)
@@ -261,6 +282,54 @@ class SimpleWheelTest(unittest.TestCase):
                 for l in string.ascii_lowercase),
             'skxqlhcnwarvgmebjptyfdzuio')
 
+    def test_set_position(self):
+        wheel_3 = Wheel(wheel_iii_spec, wheel_iii_pegs, ring_setting=3)
+        wheel_3.set_position('a')
+        self.assertEqual(wheel_3.position, 24)
+        self.assertEqual(wheel_3.position_l, 'a')
+        self.assertEqual(wheel_3.peg_positions, [21])
+
+        wheel_3.set_position('z')
+        self.assertEqual(wheel_3.position, 23)
+        self.assertEqual(wheel_3.position_l, 'z')
+        self.assertEqual(wheel_3.peg_positions, [22])
+
+        wheel_3.set_position(26)
+        self.assertEqual(wheel_3.position, 23)
+        self.assertEqual(wheel_3.position_l, 'z')
+        self.assertEqual(wheel_3.peg_positions, [22])
+
+        wheel_3.set_position(27)
+        self.assertEqual(wheel_3.position, 24)
+        self.assertEqual(wheel_3.position_l, 'a')
+        self.assertEqual(wheel_3.peg_positions, [21])
+
+        wheel_3.set_position('f')
+        self.assertEqual(wheel_3.position, 3)
+        self.assertEqual(wheel_3.position_l, 'f')
+        self.assertEqual(wheel_3.peg_positions, [16])
+
+        wheel_3.set_position(6)
+        self.assertEqual(wheel_3.position, 3)
+        self.assertEqual(wheel_3.position_l, 'f')
+        self.assertEqual(wheel_3.peg_positions, [16])
+
+        wheel_3.advance()
+        self.assertEqual(wheel_3.position, 4)
+        self.assertEqual(wheel_3.position_l, 'g')
+        self.assertEqual(wheel_3.peg_positions, [15])
+
+        wheel_3.set_position(12)
+        self.assertEqual(wheel_3.position, 9)
+        self.assertEqual(wheel_3.position_l, 'l')
+        self.assertEqual(wheel_3.peg_positions, [10])
+
+        wheel_3.advance()
+        self.assertEqual(wheel_3.position, 10)
+        self.assertEqual(wheel_3.position_l, 'm')
+        self.assertEqual(wheel_3.peg_positions, [9])
+
+
 class EnigmaTest(unittest.TestCase):
 
     def setUp(self):
@@ -281,6 +350,14 @@ class EnigmaTest(unittest.TestCase):
                 6, 20, 24, 
                 'ua pf rq so ni ey bg hl tx zj')
 
+        # Settings for Bletchley Park outreach department's Enigma
+        self.enigma_bp = Enigma(reflector_b_spec, 
+                 wheel_i_spec, wheel_i_pegs,
+                 wheel_iii_spec, wheel_iii_pegs,
+                 wheel_ii_spec, wheel_ii_pegs,
+                 1, 26, 26, 
+                 'qm we ro tu zj ps dl fg')
+
 
     def test_middle_advance(self):
         self.enigma.set_wheels('a', 'a', 't')