bgneal@1: # Copyright (C) 2012 by Brian Neal.
bgneal@1: # This file is part of Py-Enigma, the Enigma Machine simulation.
bgneal@1: # Py-Enigma is released under the MIT License (see License.txt).
bgneal@1: 
bgneal@4: """test_rotor.py - Unit tests for the Rotor class for the Enigma simulation."""
bgneal@4: 
bgneal@1: import unittest
bgneal@2: import collections
bgneal@1: import string
bgneal@1: 
bgneal@6: from ..rotors.rotor import Rotor, ALPHA_LABELS, NUMERIC_LABELS
bgneal@6: from ..rotors import RotorError
bgneal@1: 
bgneal@1: 
bgneal@1: WIRING = 'EKMFLGDQVZNTOWYHXUSPAIBRCJ'
bgneal@1: 
bgneal@1: 
bgneal@1: class SimpleRotorTestCase(unittest.TestCase):
bgneal@1:     """Basic tests to verify Rotor functionality"""
bgneal@1: 
bgneal@1:     def test_bad_wiring(self):
bgneal@1: 
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', '')
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', 'ABC')
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', '123')
bgneal@1: 
bgneal@1:         w = string.punctuation[:26]
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', w)
bgneal@1: 
bgneal@1:         w = 'ABCD' * 7
bgneal@1:         self.assertRaises(RotorError, Rotor, 'III', w[:26])
bgneal@1: 
bgneal@1:     def test_bad_ring_setting(self):
bgneal@1: 
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', WIRING, ring_setting=-1)
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', WIRING, ring_setting=26)
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', WIRING, ring_setting='A')
bgneal@1:         self.assertRaises(RotorError, Rotor, 'I', WIRING, ring_setting=None)
bgneal@1: 
bgneal@1:     def test_bad_stepping(self):
bgneal@1: 
bgneal@1:         for alpha in True, False:
bgneal@1:             self.assertRaises(RotorError, Rotor, 'I', WIRING,
bgneal@1:                     alpha_labels=alpha, stepping="0")
bgneal@1:             self.assertRaises(RotorError, Rotor, 'I', WIRING,
bgneal@1:                     alpha_labels=alpha, stepping="A0")
bgneal@1:             self.assertRaises(RotorError, Rotor, 'I', WIRING,
bgneal@1:                     alpha_labels=alpha, stepping=[1])
bgneal@1:             self.assertRaises(RotorError, Rotor, 'I', WIRING,
bgneal@1:                     alpha_labels=alpha, stepping=['A', '%', '14'])
bgneal@1:             self.assertRaises(RotorError, Rotor, 'I', WIRING,
bgneal@1:                     alpha_labels=alpha, stepping=('A', '%', '14'))
bgneal@1:     
bgneal@1:     def test_alpha_display(self):
bgneal@1: 
bgneal@1:         for r in range(26):
bgneal@1:             rotor = Rotor('I', WIRING, ring_setting=r, alpha_labels=True)
bgneal@1:             for s in ALPHA_LABELS:
bgneal@1:                 rotor.set_display(s)
bgneal@1:                 self.assertEqual(s, rotor.get_display())
bgneal@1: 
bgneal@1:     def test_numeric_display(self):
bgneal@1: 
bgneal@1:         for r in range(26):
bgneal@1:             rotor = Rotor('I', WIRING, ring_setting=r, alpha_labels=False)
bgneal@1:             for s in NUMERIC_LABELS:
bgneal@1:                 rotor.set_display(s)
bgneal@1:                 self.assertEqual(s, rotor.get_display())
bgneal@1: 
bgneal@1:     def test_wiring(self):
bgneal@5:         """Loop through all ring settings & rotor positions and test the
bgneal@5:         wiring.
bgneal@5:         
bgneal@5:         """
bgneal@3:         for r in range(26):
bgneal@3:             rotor = Rotor('I', WIRING, ring_setting=r, alpha_labels=True)
bgneal@1: 
bgneal@3:             for n, d in enumerate(ALPHA_LABELS):
bgneal@3:                 rotor.set_display(d)
bgneal@1: 
bgneal@3:                 wiring = collections.deque(WIRING)
bgneal@3:                 wiring.rotate(r - n)
bgneal@2: 
bgneal@3:                 for i in range(26):
bgneal@3:                     output = rotor.signal_in(i)
bgneal@2: 
bgneal@3:                     expected = (ord(wiring[i]) - ord('A') + r - n) % 26
bgneal@3:                     self.assertEqual(output, expected)
bgneal@3: 
bgneal@3:                     output = rotor.signal_out(expected)
bgneal@3:                     self.assertEqual(output, i)