From a311df9c2d591ca3db33b42ce50de37968f4f5b8 Mon Sep 17 00:00:00 2001
From: James Kasten <jdkasten@umich.edu>
Date: Mon, 9 Feb 2015 02:46:29 -0800
Subject: [PATCH] Add tests for enhancments UI

---
 .../client/tests/display/enhancements_test.py | 58 +++++++++++++++++++
 1 file changed, 58 insertions(+)
 create mode 100644 letsencrypt/client/tests/display/enhancements_test.py

diff --git a/letsencrypt/client/tests/display/enhancements_test.py b/letsencrypt/client/tests/display/enhancements_test.py
new file mode 100644
index 000000000..dfdac52e2
--- /dev/null
+++ b/letsencrypt/client/tests/display/enhancements_test.py
@@ -0,0 +1,58 @@
+"""Module for enhancement UI."""
+import logging
+import unittest
+
+import mock
+
+from letsencrypt.client import errors
+from letsencrypt.client.display import display_util
+
+
+class AskTest(unittest.TestCase):
+    """Test the ask method."""
+    def setUp(self):
+        logging.disable(logging.CRITICAL)
+
+    def tearDown(self):
+        logging.disable(logging.NOTSET)
+
+    @classmethod
+    def _call(cls, enhancement):
+        from letsencrypt.client.display.enhancements import ask
+        return ask(enhancement)
+
+    @mock.patch("letsencrypt.client.display.enhancements.util")
+    def test_redirect(self, mock_util):
+        mock_util().menu.return_value = (display_util.OK, 1)
+        self.assertTrue(self._call("redirect"))
+
+    def test_key_error(self):
+        self.assertRaises(
+            errors.LetsEncryptClientError, self._call, "unknown_enhancement")
+
+
+class RedirectTest(unittest.TestCase):
+    """Test the redirect_by_default method."""
+    @classmethod
+    def _call(cls):
+        from letsencrypt.client.display.enhancements import redirect_by_default
+        return redirect_by_default()
+
+    @mock.patch("letsencrypt.client.display.enhancements.util")
+    def test_secure(self, mock_util):
+        mock_util().menu.return_value = (display_util.OK, 1)
+        self.assertTrue(self._call())
+
+    @mock.patch("letsencrypt.client.display.enhancements.util")
+    def test_cancel(self, mock_util):
+        mock_util().menu.return_value = (display_util.CANCEL, 1)
+        self.assertFalse(self._call())
+
+    @mock.patch("letsencrypt.client.display.enhancements.util")
+    def test_easy(self, mock_util):
+        mock_util().menu.return_value = (display_util.OK, 0)
+        self.assertFalse(self._call())
+
+
+if __name__ == "__main__":
+    unittest.main()