diff --git a/python/src/uagents/agent.py b/python/src/uagents/agent.py index 425fd92d..9c4656dc 100644 --- a/python/src/uagents/agent.py +++ b/python/src/uagents/agent.py @@ -1377,6 +1377,7 @@ def __init__( registration_policy: Optional[BatchRegistrationPolicy] = None, ledger: Optional[LedgerClient] = None, wallet: Optional[LocalWallet] = None, + seed: Optional[str] = None, test: bool = True, loop: Optional[asyncio.AbstractEventLoop] = None, log_level: Union[int, str] = logging.INFO, @@ -1390,7 +1391,8 @@ def __init__( endpoint (Optional[Union[str, List[str], Dict[str, dict]]]): The endpoint configuration. agentverse (Optional[Union[str, Dict[str, str]]]): The agentverse configuration. registration_policy (Optional[BatchRegistrationPolicy]): The registration policy. - wallet (Optional[LocalWallet]): The wallet for the bureau. + wallet (Optional[LocalWallet]): The wallet for the bureau (overrides 'seed'). + seed (Optional[str]): The seed phrase for the wallet (overridden by 'wallet'). test (Optional[bool]): True if the bureau will register and transact on the testnet. loop (Optional[asyncio.AbstractEventLoop]): The event loop. log_level (Union[int, str]): The logging level for the bureau. @@ -1412,6 +1414,16 @@ def __init__( almanac_api_url = f"{self._agentverse['http_prefix']}://{self._agentverse['base_url']}/v1/almanac" almanac_contract = get_almanac_contract(test) + if wallet and seed: + self._logger.warning( + "Ignoring 'seed' argument because 'wallet' is provided." + ) + elif seed: + wallet = LocalWallet( + PrivateKey(derive_key_from_seed(seed, LEDGER_PREFIX, 0)), + prefix=LEDGER_PREFIX, + ) + if registration_policy is not None: if ( isinstance(registration_policy, BatchLedgerRegistrationPolicy) diff --git a/python/tests/test_bureau.py b/python/tests/test_bureau.py new file mode 100644 index 00000000..84119c01 --- /dev/null +++ b/python/tests/test_bureau.py @@ -0,0 +1,105 @@ +from cosmpy.aerial.wallet import LocalWallet + +from uagents import Agent, Bureau +from uagents.registration import ( + AgentEndpoint, + BatchLedgerRegistrationPolicy, + DefaultBatchRegistrationPolicy, + DefaultRegistrationPolicy, + LedgerBasedRegistrationPolicy, +) + +ALICE_ENDPOINT = AgentEndpoint(url="http://alice:8000/submit", weight=1) +BOB_ENDPOINT = AgentEndpoint(url="http://bob:8000/submit", weight=1) +BUREAU_ENDPOINT = AgentEndpoint(url="http://bureau:8000/submit", weight=1) + + +bureau_wallet = LocalWallet.generate() + + +def test_bureau_updates_agents_no_ledger_batch(): + alice = Agent(name="alice", endpoint=ALICE_ENDPOINT.url) + bob = Agent(name="bob", endpoint=BOB_ENDPOINT.url) + + assert alice._endpoints == [ALICE_ENDPOINT] + assert bob._endpoints == [BOB_ENDPOINT] + + assert isinstance(alice._registration_policy, DefaultRegistrationPolicy) + assert isinstance(bob._registration_policy, DefaultRegistrationPolicy) + + bureau = Bureau(agents=[alice, bob], endpoint=BUREAU_ENDPOINT.url) + + assert alice._endpoints == [BUREAU_ENDPOINT] + assert bob._endpoints == [BUREAU_ENDPOINT] + + assert isinstance(bureau._registration_policy, DefaultBatchRegistrationPolicy) + assert bureau._registration_policy._ledger_policy is None + assert isinstance(alice._registration_policy, LedgerBasedRegistrationPolicy) + assert isinstance(bob._registration_policy, LedgerBasedRegistrationPolicy) + + +def test_bureau_updates_agents_with_wallet(): + alice = Agent(name="alice", endpoint=ALICE_ENDPOINT.url) + bob = Agent(name="bob", endpoint=BOB_ENDPOINT.url) + + assert isinstance(alice._registration_policy, DefaultRegistrationPolicy) + assert isinstance(bob._registration_policy, DefaultRegistrationPolicy) + + bureau = Bureau(agents=[alice, bob], wallet=bureau_wallet) + + assert alice._endpoints == [] + assert bob._endpoints == [] + + assert isinstance(bureau._registration_policy, DefaultBatchRegistrationPolicy) + assert isinstance( + bureau._registration_policy._ledger_policy, BatchLedgerRegistrationPolicy + ) + assert alice._registration_policy is None + assert bob._registration_policy is None + + +def test_bureau_updates_agents_with_seed(): + alice = Agent(name="alice", endpoint=ALICE_ENDPOINT.url) + bob = Agent(name="bob", endpoint=BOB_ENDPOINT.url) + + assert isinstance(alice._registration_policy, DefaultRegistrationPolicy) + assert isinstance(bob._registration_policy, DefaultRegistrationPolicy) + + bureau = Bureau( + agents=[alice, bob], + endpoint=BUREAU_ENDPOINT.url, + seed="bureau test seed phrase", + ) + + assert isinstance(bureau._registration_policy, DefaultBatchRegistrationPolicy) + assert isinstance( + bureau._registration_policy._ledger_policy, BatchLedgerRegistrationPolicy + ) + assert alice._registration_policy is None + assert bob._registration_policy is None + + +def test_bureau_updates_agents_wallet_overrides_seed(): + alice = Agent(name="alice", endpoint=ALICE_ENDPOINT.url) + bob = Agent(name="bob", endpoint=BOB_ENDPOINT.url) + + assert isinstance(alice._registration_policy, DefaultRegistrationPolicy) + assert isinstance(bob._registration_policy, DefaultRegistrationPolicy) + + bureau = Bureau( + agents=[alice, bob], + endpoint=BUREAU_ENDPOINT.url, + wallet=bureau_wallet, + seed="bureau test seed phrase", + ) + + assert isinstance(bureau._registration_policy, DefaultBatchRegistrationPolicy) + assert isinstance( + bureau._registration_policy._ledger_policy, BatchLedgerRegistrationPolicy + ) + assert ( + bureau._registration_policy._ledger_policy._wallet.address() + == bureau_wallet.address() + ) + assert alice._registration_policy is None + assert bob._registration_policy is None