Test encrypting floats.

Quick refactor of cli testing.

phe/tests/cli_test.py

@@ -1,18 +1,18 @@
 import json
 from unittest import TestCase
 import tempfile
-from contextlib import redirect_stdout, redirect_stderr
 
 import io
 
 import sys
 import click
 from click.testing import CliRunner
 
+import phe.command_line
 from phe.command_line import cli
 
 
-class TestConsole(TestCase):
+class TestConsoleBasics(TestCase):
 
     def test_cli_includes_help(self):
         runner = CliRunner()
@@ -69,34 +69,53 @@ def test_extract_public_key(self):
                 assert '"mu":' not in written_data
                 assert '"lamdba":' not in written_data
 
-    def test_encrypt_positive_integers(self):
-        """Test encrypting an integer"""
-        runner = CliRunner()
 
-        numbers = [0, 1, 2, 5, 10, '1', '10550']
+class TestConsoleEncryption(TestCase):
 
-        with tempfile.NamedTemporaryFile() as private_keyfile:
-            with tempfile.NamedTemporaryFile() as public_keyfile:
-                runner.invoke(cli, ['generate', '--keysize', '256', private_keyfile.name])
-                runner.invoke(cli, ['extract', private_keyfile.name, public_keyfile.name])
+    @classmethod
+    def setUpClass(cls):
+        """Generate a keypair and extract the public key.
+        """
+        cls.private_keyfile = tempfile.NamedTemporaryFile()
+        cls.public_keyfile = tempfile.NamedTemporaryFile()
 
-                for num in numbers:
-                    result = runner.invoke(cli, ['encrypt', public_keyfile.name, str(num)])
-                    assert result.exit_code == 0
+        cls.runner = CliRunner()
+        cls.runner.invoke(cli, ['generate', '--keysize', '256', cls.private_keyfile.name])
+        cls.runner.invoke(cli, ['extract', cls.private_keyfile.name, cls.public_keyfile.name])
 
+    @classmethod
+    def tearDownClass(cls):
+        cls.private_keyfile.close()
+        cls.public_keyfile.close()
 
-    def test_encrypt_signed_integers(self):
-        """Test encrypting an integer"""
-        runner = CliRunner()
+    def test_encrypt_positive_integers(self):
+        numbers = [0, 1, 2, 5, 10, '1', '10550']
+
+        for num in numbers:
+            result = self.runner.invoke(cli, ['encrypt', self.public_keyfile.name, str(num)])
+            assert result.exit_code == 0
 
+    def test_encrypt_signed_integers(self):
+        """encrypting positive and negative integer"""
         numbers = [0, 1, -1, 10, '1', '-10550']
 
-        with tempfile.NamedTemporaryFile() as private_keyfile:
-            with tempfile.NamedTemporaryFile() as public_keyfile:
-                runner.invoke(cli, ['generate', '--keysize', '256', private_keyfile.name])
-                runner.invoke(cli, ['extract', private_keyfile.name, public_keyfile.name])
+        for num in numbers:
+            result = self.runner.invoke(cli, ['encrypt', self.public_keyfile.name, "--", str(num)])
+            assert result.exit_code == 0
+
+    def test_encrypt_float(self):
+        numbers = [0.0, 1.1, -0.0001, 100000.01, '1e-20', '-10550e20']
 
-                for num in numbers:
-                    result = runner.invoke(cli, ['encrypt', public_keyfile.name, "--", str(num)])
-                    assert result.exit_code == 0
+
+        for num in numbers:
+            result = self.runner.invoke(cli, ['encrypt', self.public_keyfile.name, "--", str(num)])
+            assert result.exit_code == 0
+
+    def test_encrypt_to_stdout(self):
+        """Test encrypting and writing output to a file"""
+        numbers = [0.0, 1.1, -0.0001, 100000.01, '1e-20', '-10550e20']
+
+        for num in numbers:
+            result = self.runner.invoke(cli, ['encrypt', self.public_keyfile.name, "--", str(num)])
+            assert result.exit_code == 0