# HG changeset patch
# User Brian Neal <bgneal@gmail.com>
# Date 1372816830 18000
# Node ID 96fe2240936d91f4ab465acf26dc8d6b0584f2fb
# Parent  dae8d52023e7c4775db3d2dad294a33ede3b85bd
Better test for encrypt padding. Fix bug in padding that test revealed.

diff -r dae8d52023e7 -r 96fe2240936d m209/procedure.py
--- a/m209/procedure.py	Sun Jun 30 20:26:32 2013 -0500
+++ b/m209/procedure.py	Tue Jul 02 21:00:30 2013 -0500
@@ -156,8 +156,11 @@
         i = ciphertext.rfind(' ')
         if i != -1:
             x_count = 6 - (len(ciphertext) - i)
-            if x_count:
-                ciphertext = ciphertext + 'X' * x_count
+        else:   # only 1 group
+            x_count = 5 - len(ciphertext)
+
+        if x_count:
+            ciphertext = ciphertext + 'X' * x_count
 
         # Add the message indicators to pad each end of the message
 
diff -r dae8d52023e7 -r 96fe2240936d m209/tests/test_procedure.py
--- a/m209/tests/test_procedure.py	Sun Jun 30 20:26:32 2013 -0500
+++ b/m209/tests/test_procedure.py	Tue Jul 02 21:00:30 2013 -0500
@@ -46,11 +46,14 @@
         self.assertEqual(plaintext[:len(PLAINTEXT)], PLAINTEXT)
 
     def test_encrypt_padding(self):
-        pt = 'THE PIZZA HAS ARRIVED STOP WILL PLAY DAYZ AFTER SUPPER STOP'
-        ct = self.proc.encrypt(pt)
-        groups = ct.split()
-        counts_good = [len(group) == 5 for group in groups]
-        self.assertTrue(all(counts_good))
+        """Ensure we pad the final group out to 5 chars."""
+
+        for n in range(1, 21):
+            pt = 'A' * n
+            ct = self.proc.encrypt(pt)
+            groups = ct.split()
+            counts_good = [len(group) == 5 for group in groups]
+            self.assertTrue(all(counts_good), "Failed on n = %s; ct = %s" % (n, ct))
 
     def test_blair_decrypt(self):