test: add bureau registration test cases

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)

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