# HG changeset patch
# User Brian Neal <bgneal@gmail.com>
# Date 1338173627 18000
# Node ID f4051bcac8c5197d1963cba967a549cc75bd4735
# Parent  ce384011521446f9e1a3fdf14ab72b4baa3b7396
Fix for the Rotor.notch_over_pawl() bug.

Added 2 more tests for rotors: test_notches & test_rotate.

diff -r ce3840115214 -r f4051bcac8c5 enigma/rotors/rotor.py
--- a/enigma/rotors/rotor.py	Sun May 27 21:20:34 2012 -0500
+++ b/enigma/rotors/rotor.py	Sun May 27 21:53:47 2012 -0500
@@ -125,22 +125,27 @@
         for i, v in enumerate(self.entry_map):
             self.exit_map[v] = i
 
-        # configure ring labels
-        self.ring_labels = ALPHA_LABELS
+        # build a map of display values to positions
+        self.display_map = {}
+        for n in range(26):
+            self.display_map[chr(ord('A') + n)] = (n - self.ring_setting) % 26
 
-        # for mapping window display values to indices
-        self.display_map = dict(zip(self.ring_labels, range(26)))
+        # build a reverse map of position mapped to display values
+        self.pos_map = {v : k for k, v in self.display_map.items()}
 
-        # build step list: this is a list of positions where our notches are in
+        # build step set; this is a set of positions where our notches are in
         # place to allow the pawls to move
         self.step_set = set()
         if stepping is not None:
             for pos in stepping:
                 if pos in self.display_map:
-                    self.step_set.add(self.display_map[pos])
+                    self.step_set.add(pos)
                 else:
                     raise RotorError("stepping: %s" % pos)
 
+        # initialize our position and display value:
+        self.set_display('A')
+
     def set_display(self, val):
         """Spin the rotor such that the string val appears in the operator
         window.
@@ -158,18 +163,15 @@
         """
         s = val.upper()
         if s not in self.display_map:
-            raise RotorError("bad display value %s" % s)
+            raise RotorError("bad display value %s" % val)
 
-        index = self.display_map[s]
-
-        self.pos = (index - self.ring_setting) % 26
+        self.pos = self.display_map[s]
+        self.display_val = s
         self.rotations = 0
 
     def get_display(self):
         """Returns what is currently being displayed in the operator window."""
-
-        index = (self.pos + self.ring_setting) % 26
-        return self.ring_labels[index]
+        return self.display_val
 
     def signal_in(self, n):
         """Simulate a signal entering the rotor from the right at a given pin
@@ -212,10 +214,11 @@
         False otherwise.
 
         """
-        return self.pos in self.step_set
+        return self.display_val in self.step_set
 
     def rotate(self):
         """Rotate the rotor forward due to mechanical stepping action."""
 
         self.pos = (self.pos + 1) % 26
+        self.display_val = self.pos_map[self.pos]
         self.rotations += 1
diff -r ce3840115214 -r f4051bcac8c5 enigma/tests/test_rotor.py
--- a/enigma/tests/test_rotor.py	Sun May 27 21:20:34 2012 -0500
+++ b/enigma/tests/test_rotor.py	Sun May 27 21:53:47 2012 -0500
@@ -10,6 +10,8 @@
 
 from ..rotors.rotor import Rotor, ALPHA_LABELS
 from ..rotors import RotorError
+from ..rotors.data import ROTORS
+from ..rotors.factory import create_rotor
 
 
 WIRING = 'EKMFLGDQVZNTOWYHXUSPAIBRCJ'
@@ -75,3 +77,35 @@
 
                     output = rotor.signal_out(expected)
                     self.assertEqual(output, i)
+
+    def test_notches(self):
+        """For every rotor we simulate, ensure that the notch setting is correct
+        regardless of the ring setting.
+
+        """
+        rotor_names = ROTORS.keys()
+        for rotor_name in rotor_names:
+            notches = ROTORS[rotor_name]['stepping']
+            if notches is None:
+                continue
+            for r in range(26):
+                rotor = create_rotor(rotor_name, r)
+                rotor.set_display('A')
+
+                for n in range(26):
+                    over_notch = rotor.get_display() in notches
+                    self.assertEqual(rotor.notch_over_pawl(), over_notch)
+                    rotor.rotate()
+
+    def test_rotate(self):
+
+        for r in range(26):
+            rotor1 = Rotor('X', WIRING, ring_setting=r)
+            rotor2 = Rotor('Y', WIRING, ring_setting=r)
+
+            rotor2.set_display('A')
+
+            for i, d in enumerate(ALPHA_LABELS):
+                rotor1.set_display(d)
+                self.assertEqual(rotor1.get_display(), rotor2.get_display())
+                rotor2.rotate()