diff --git a/phe/tests/cli_test.py b/phe/tests/cli_test.py index ead5118..895d3c6 100644 --- a/phe/tests/cli_test.py +++ b/phe/tests/cli_test.py @@ -1,7 +1,6 @@ import json from unittest import TestCase import tempfile -from contextlib import redirect_stdout, redirect_stderr import io @@ -9,10 +8,11 @@ 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