add some util functions to jsonrpc client, testnet and identifier.Intent

Author: Xiao Li

.gitignore

@@ -3,7 +3,7 @@
 /build
 /dist
 __pycache__
-/src/diem_client_sdk.egg-info
+/src/*.egg-info
 
 # JetBrains IDEs configuration
 .idea

Makefile

@@ -23,7 +23,7 @@ install:
 	./venv/bin/python setup.py develop
 
 test: format install
-	./venv/bin/pytest tests/test_* examples/* -k "$(TEST)"
+	./venv/bin/pytest tests/test_* examples/* -k "$(TEST)" -vv
 
 cover: install
 	./venv/bin/pytest --cov-report html --cov=src tests

src/diem/identifier/__init__.py

@@ -57,16 +57,22 @@ def __init__(
         sub_address: typing.Optional[bytes],
         currency_code: str,
         amount: int,
+        hrp: str,
     ) -> None:
         self.account_address = account_address
         self.sub_address = sub_address
         self.currency_code = currency_code
         self.amount = amount
+        self.hrp = hrp
 
     @property
     def account_address_bytes(self) -> bytes:
         return self.account_address.to_bytes()
 
+    @property
+    def account_id(self) -> str:
+        return encode_account(self.account_address, self.sub_address, self.hrp)
+
 
 def encode_intent(encoded_account_identifier: str, currency_code: str, amount: int) -> str:
     """
@@ -109,6 +115,7 @@ def decode_intent(encoded_intent_identifier: str, hrp: str) -> Intent:
         sub_address=sub_address,
         currency_code=currency_code,
         amount=amount,
+        hrp=hrp,
     )
 
 

src/diem/jsonrpc/client.py

@@ -11,6 +11,7 @@
 import typing
 import random, queue
 from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
 
 from .. import diem_types, utils
 from . import jsonrpc_pb2 as rpc
@@ -55,7 +56,7 @@ class WaitForTransactionTimeout(Exception):
     pass
 
 
-class AccountNotFoundError(Exception):
+class AccountNotFoundError(ValueError):
     pass
 
 
@@ -205,10 +206,7 @@ def get_parent_vasp_account(
         could not find the account by the parent_vasp_address found in ChildVASP account.
         """
 
-        account = self.get_account(vasp_account_address)
-        if account is None:
-            hex = utils.account_address_hex(vasp_account_address)
-            raise AccountNotFoundError(f"account not found by address: {hex}")
+        account = self.must_get_account(vasp_account_address)
 
         if account.role.type == constants.ACCOUNT_ROLE_PARENT_VASP:
             return account
@@ -218,6 +216,33 @@ def get_parent_vasp_account(
         hex = utils.account_address_hex(vasp_account_address)
         raise ValueError(f"given account address({hex}) is not a VASP account: {account}")
 
+    def get_base_url_and_compliance_key(
+        self, account_address: typing.Union[diem_types.AccountAddress, str]
+    ) -> typing.Tuple[str, Ed25519PublicKey]:
+        """get base_url and compliance key
+
+        ParentVASP or Designated Dealer account role has base_url and compliance key setup, which
+        are used for offchain API communication.
+        """
+        account = self.must_get_account(account_address)
+
+        if account.role.compliance_key and account.role.base_url:
+            key = Ed25519PublicKey.from_public_bytes(bytes.fromhex(account.role.compliance_key))
+            return (account.role.base_url, key)
+        if account.role.parent_vasp_address:
+            return self.get_base_url_and_compliance_key(account.role.parent_vasp_address)
+
+        raise ValueError(f"could not find base_url and compliance_key from account: {account}")
+
+    def must_get_account(self, account_address: typing.Union[diem_types.AccountAddress, str]) -> rpc.Account:
+        """must_get_account raises AccountNotFoundError if account could not be found by given address"""
+
+        account = self.get_account(account_address)
+        if account is None:
+            hex = utils.account_address_hex(account_address)
+            raise AccountNotFoundError(f"account not found by address: {hex}")
+        return account
+
     def get_account_sequence(self, account_address: typing.Union[diem_types.AccountAddress, str]) -> int:
         """get on-chain account sequence number
 

src/diem/local_account.py

@@ -35,9 +35,11 @@ def generate() -> "LocalAccount":
         return LocalAccount(private_key)
 
     private_key: Ed25519PrivateKey
+    compliance_key: Ed25519PrivateKey
 
     def __init__(self, private_key: Ed25519PrivateKey) -> None:
         self.private_key = private_key
+        self.compliance_key = Ed25519PrivateKey.generate()
 
     @property
     def auth_key(self) -> AuthKey:
@@ -55,6 +57,10 @@ def public_key_bytes(self) -> bytes:
     def public_key(self) -> Ed25519PublicKey:
         return self.private_key.public_key()
 
+    @property
+    def compliance_public_key_bytes(self) -> bytes:
+        return utils.public_key_bytes(self.compliance_key.public_key())
+
     def sign(self, txn: diem_types.RawTransaction) -> diem_types.SignedTransaction:
         """Create signed transaction for given raw transaction"""
 

src/diem/testnet.py

@@ -21,9 +21,9 @@
 """
 
 import requests
-import typing
+import typing, time
 
-from . import diem_types, jsonrpc, utils, local_account, serde_types, auth_key, chain_ids, lcs
+from . import diem_types, jsonrpc, utils, local_account, serde_types, auth_key, chain_ids, lcs, stdlib, LocalAccount
 
 
 JSON_RPC_URL: str = "https://testnet.diem.com/v1"
@@ -40,6 +40,62 @@ def create_client() -> jsonrpc.Client:
     return jsonrpc.Client(JSON_RPC_URL)
 
 
+def gen_vasp_account(base_url: str) -> LocalAccount:
+    account = gen_account()
+    exec_txn(
+        account,
+        stdlib.encode_rotate_dual_attestation_info_script(
+            new_url=base_url.encode("utf-8"), new_key=account.compliance_public_key_bytes
+        ),
+    )
+    return account
+
+
+def gen_account() -> LocalAccount:
+    """generates a Testnet onchain account"""
+
+    return Faucet(create_client()).gen_account()
+
+
+def gen_child_vasp(
+    parent_vasp: LocalAccount, initial_balance: int = 10_000_000_000, currency: str = TEST_CURRENCY_CODE
+) -> LocalAccount:
+    child_vasp = LocalAccount.generate()
+    exec_txn(
+        parent_vasp,
+        stdlib.encode_create_child_vasp_account_script(
+            coin_type=utils.currency_code(currency),
+            child_address=child_vasp.account_address,
+            auth_key_prefix=child_vasp.auth_key.prefix(),
+            add_all_currencies=False,
+            child_initial_balance=initial_balance,
+        ),
+    )
+    return child_vasp
+
+
+def exec_txn(sender: LocalAccount, script: diem_types.Script) -> jsonrpc.Transaction:
+    """create, submit and wait for the transaction"""
+
+    client = create_client()
+    sender_account_sequence = client.get_account_sequence(sender.account_address)
+    expire = 30
+    txn = sender.sign(
+        diem_types.RawTransaction(  # pyre-ignore
+            sender=sender.account_address,
+            sequence_number=diem_types.st.uint64(sender_account_sequence),
+            payload=diem_types.TransactionPayload__Script(value=script),
+            max_gas_amount=1_000_000,
+            gas_unit_price=0,
+            gas_currency_code=TEST_CURRENCY_CODE,
+            expiration_timestamp_secs=int(time.time()) + expire,
+            chain_id=CHAIN_ID,
+        )
+    )
+    client.submit(txn)
+    return client.wait_for_transaction(txn, timeout_secs=expire)
+
+
 class Faucet:
     """Faucet service is a proxy server to mint coins for your test account on Testnet
 
@@ -57,9 +113,9 @@ def __init__(
         self._retry: jsonrpc.Retry = retry or jsonrpc.Retry(5, 0.2, Exception)
         self._session: requests.Session = requests.Session()
 
-    def gen_account(self, currency_code: str = TEST_CURRENCY_CODE) -> local_account.LocalAccount:
-        account = local_account.LocalAccount.generate()
-        self.mint(account.auth_key.hex(), 5_000_000_000, currency_code)
+    def gen_account(self, currency_code: str = TEST_CURRENCY_CODE) -> LocalAccount:
+        account = LocalAccount.generate()
+        self.mint(account.auth_key.hex(), 100_000_000_000, currency_code)
         return account
 
     def mint(self, authkey: str, amount: int, currency_code: str) -> None:

src/diem/utils.py

@@ -155,3 +155,10 @@ def decode_transaction_script(
         return decode_transaction_script(txn.script_bytes)
 
     raise TypeError(f"unknown transaction type: {txn}")
+
+
+def balance(account: jsonrpc.Account, currency: str) -> int:
+    for b in account.balances:
+        if b.currency == currency:
+            return b.amount
+    return 0

tests/test_identifier.py

@@ -129,6 +129,8 @@ def test_intent_identifier():
     assert intent.currency_code == "Coin1"
     assert intent.amount == 123
 
+    assert account_id == intent.account_id
+
 
 def test_intent_identifier_with_sub_address():
     account_id = identifier.encode_account(test_onchain_address, test_sub_address, "lbr")

tests/test_testnet.py

@@ -346,6 +346,42 @@ def test_get_parent_vasp_account_with_non_vasp_account_address():
         client.get_parent_vasp_account(utils.TREASURY_ADDRESS)
 
 
+def test_gen_vasp_account():
+    account = testnet.gen_vasp_account("http://hello.com")
+    child_vasp = testnet.gen_child_vasp(account)
+
+    client = testnet.create_client()
+    assert client.get_account(account.account_address).role.type == "parent_vasp"
+    assert client.get_account(child_vasp.account_address).role.type == "child_vasp"
+
+
+def test_get_base_url_and_compliance_key():
+    client = testnet.create_client()
+
+    parent_vasp = testnet.gen_vasp_account("http://hello.com")
+    child_vasp = testnet.gen_child_vasp(parent_vasp)
+
+    base_url, key = client.get_base_url_and_compliance_key(child_vasp.account_address)
+    assert base_url == "http://hello.com"
+    assert utils.public_key_bytes(key) == parent_vasp.compliance_public_key_bytes
+    base_url, key = client.get_base_url_and_compliance_key(parent_vasp.account_address)
+    assert base_url == "http://hello.com"
+    assert utils.public_key_bytes(key) == parent_vasp.compliance_public_key_bytes
+
+
+def test_account_not_found_error_when_get_base_url_and_compliance_key_for_invalid_account():
+    client = testnet.create_client()
+    account = LocalAccount.generate()
+    with pytest.raises(jsonrpc.AccountNotFoundError):
+        client.get_base_url_and_compliance_key(account.account_address)
+
+
+def test_value_error_when_get_base_url_and_compliance_key_for_account_has_no_base_url():
+    client = testnet.create_client()
+    with pytest.raises(ValueError):
+        client.get_base_url_and_compliance_key(utils.TREASURY_ADDRESS)
+
+
 def create_child_vasp_txn(
     parent_vasp: LocalAccount, child_vasp: LocalAccount, seq: int = 0
 ) -> diem_types.RawTransaction:

tests/test_utils.py

@@ -71,3 +71,15 @@ def test_decode_transaction_script():
 
     with pytest.raises(TypeError):
         utils.decode_transaction_script(False)
+
+
+def test_balance():
+    account = jsonrpc.Account(
+        balances=[
+            jsonrpc.Amount(amount=32, currency="Coin1"),
+            jsonrpc.Amount(amount=33, currency="Coin2"),
+        ]
+    )
+    assert utils.balance(account, "Coin1") == 32
+    assert utils.balance(account, "Coin2") == 33
+    assert utils.balance(account, "unknown") == 0