diff --git a/squeaknode/core/profiles.py b/squeaknode/core/profiles.py index 1afb5c5c..ed78d647 100644 --- a/squeaknode/core/profiles.py +++ b/squeaknode/core/profiles.py @@ -68,3 +68,11 @@ def validate_profile_name(profile_name: str) -> None: raise Exception( "Profile name cannot be empty.", ) + + +def get_profile_private_key(profile: SqueakProfile) -> bytes: + if profile.private_key is None: + raise Exception("Profile: {} does not have a private key.".format( + profile, + )) + return profile.private_key diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py index f69ea45b..26d9882e 100644 --- a/squeaknode/node/squeak_controller.py +++ b/squeaknode/node/squeak_controller.py @@ -45,6 +45,7 @@ from squeaknode.core.peer_address import PeerAddress from squeaknode.core.peers import create_saved_peer from squeaknode.core.profiles import create_contact_profile from squeaknode.core.profiles import create_signing_profile +from squeaknode.core.profiles import get_profile_private_key from squeaknode.core.received_offer import ReceivedOffer from squeaknode.core.received_payment import ReceivedPayment from squeaknode.core.received_payment_summary import ReceivedPaymentSummary @@ -320,13 +321,9 @@ class SqueakController: profile = self.get_squeak_profile(profile_id) if profile is None: raise Exception("Profile with id: {} does not exist.".format( - profile_id + profile_id, )) - if profile.private_key is None: - raise Exception("Profile with id: {} does not have a private key.".format( - profile_id - )) - return profile.private_key + return get_profile_private_key(profile) def create_peer(self, peer_name: str, peer_address: PeerAddress): squeak_peer = create_saved_peer( diff --git a/tests/conftest.py b/tests/conftest.py index 34d6a3d5..ba944fab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -54,6 +54,11 @@ def signing_key(): yield CSigningKey.generate() +@pytest.fixture +def signing_key_bytes(signing_key): + yield str(signing_key).encode() + + @pytest.fixture def address(signing_key): verifying_key = signing_key.get_verifying_key() diff --git a/tests/core/test_profiles.py b/tests/core/test_profiles.py index 2fd0e97a..d2e7973d 100644 --- a/tests/core/test_profiles.py +++ b/tests/core/test_profiles.py @@ -23,6 +23,7 @@ import pytest from squeaknode.core.profiles import create_contact_profile from squeaknode.core.profiles import create_signing_profile +from squeaknode.core.profiles import get_profile_private_key @pytest.fixture @@ -71,3 +72,15 @@ def test_create_contact_profile_invalid_address(profile_name, invalid_address_st with pytest.raises(Exception) as excinfo: create_contact_profile(profile_name, invalid_address_str) assert "Invalid squeak address" in str(excinfo.value) + + +def test_get_profile_private_key(signing_profile, signing_key_bytes): + private_key = get_profile_private_key(signing_profile) + + assert private_key == signing_key_bytes + + +def test_get_profile_private_key_missing(contact_profile): + with pytest.raises(Exception) as excinfo: + get_profile_private_key(contact_profile) + assert "does not have a private key" in str(excinfo.value)