Skip to content

Commit ec5241a

Browse files
authored
feat: group compatible pytests to use shared cluster (#1468)
1 parent 859276b commit ec5241a

11 files changed

Lines changed: 283 additions & 218 deletions

pytest/common_lib/contract_state.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,9 @@ def from_json(resharing_data: dict):
278278
for k in resharing_data["reshared_keys"]:
279279
reshared_keys_list.append(
280280
KeyForDomain(
281-
domain_id=k["domain_id"], attempt_id=k["attempt"], key="placeholder"
281+
domain_id=k["domain_id"],
282+
attempt_id=k["attempt"],
283+
key={"placeholder": None},
282284
)
283285
)
284286

@@ -373,7 +375,9 @@ def from_json(data):
373375
for k in data["generated_keys"]:
374376
generated_keys_list.append(
375377
KeyForDomain(
376-
domain_id=k["domain_id"], attempt_id=k["attempt"], key="placeholder"
378+
domain_id=k["domain_id"],
379+
attempt_id=k["attempt"],
380+
key={"placeholder": None},
377381
)
378382
)
379383

@@ -430,8 +434,10 @@ def print(self):
430434
class ContractState:
431435
def get_running_domains(self) -> List[Domain]:
432436
if self.state == ProtocolState.RUNNING:
437+
assert isinstance(self.protocol_state, RunningProtocolState)
433438
return self.protocol_state.domains.domains
434439
elif self.state == ProtocolState.RESHARING:
440+
assert isinstance(self.protocol_state, ResharingProtocolState)
435441
return self.protocol_state.previous_running_state.domains.domains
436442

437443
assert False, "expected running state"
@@ -451,6 +457,7 @@ def __init__(self, data):
451457

452458
def keyset(self) -> Keyset | None:
453459
if self.state == ProtocolState.RUNNING:
460+
assert isinstance(self.protocol_state, RunningProtocolState)
454461
return self.protocol_state.keyset
455462
return None
456463

pytest/common_lib/shared/metrics.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,12 @@ class IntMetricName(str, Enum):
1313
"mpc_owned_num_presignatures_with_offline_participant"
1414
)
1515
MPC_INDEXER_LATEST_BLOCK_HEIGHT = "mpc_indexer_latest_block_height"
16+
MPC_PENDING_SIGNATURES_QUEUE_ATTEMPTS_GENERATED = (
17+
"mpc_pending_signatures_queue_attempts_generated"
18+
)
19+
MPC_PENDING_CKDS_QUEUE_ATTEMPTS_GENERATED = (
20+
"mpc_pending_ckds_queue_attempts_generated"
21+
)
1622

1723

1824
class DictMetricName(str, Enum):

pytest/common_lib/shared/mpc_cluster.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,12 @@
99
from common_lib import signature
1010
from common_lib import ckd
1111
from common_lib.constants import TGAS
12-
from common_lib.contract_state import ContractState, ProtocolState, SignatureScheme
12+
from common_lib.contract_state import (
13+
ContractState,
14+
ProtocolState,
15+
SignatureScheme,
16+
RunningProtocolState,
17+
)
1318
from common_lib.contracts import ContractMethod
1419
from common_lib.migration_state import (
1520
BackupServiceInfo,
@@ -209,6 +214,7 @@ def update_participant_status(self, assert_contract=True):
209214
self.nodes = nodes
210215
if assert_contract:
211216
contract_state = self.contract_state()
217+
assert isinstance(contract_state.protocol_state, RunningProtocolState)
212218
assert len(
213219
contract_state.protocol_state.parameters.participants.participants
214220
) == len(self.mpc_nodes)
@@ -274,6 +280,7 @@ def add_domains(
274280
state = self.contract_state()
275281
state.print()
276282
assert state.is_state(ProtocolState.RUNNING), "require running state"
283+
assert isinstance(state.protocol_state, RunningProtocolState)
277284
domains_to_add = []
278285
next_domain_id = state.protocol_state.next_domain_id()
279286
for scheme in schemes:

pytest/tests/shared_cluster_tests/__init__.py

Whitespace-only changes.
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import pytest
2+
import sys
3+
import pathlib
4+
import atexit
5+
6+
sys.path.append(str(pathlib.Path(__file__).resolve().parents[2]))
7+
8+
from common_lib import shared, contracts, contract_state
9+
10+
11+
@pytest.fixture(scope="package")
12+
def shared_cluster():
13+
"""
14+
Spins up a cluster with three nodes, initializes the contract and adds domains. Returns the cluster in a running state.
15+
"""
16+
cluster, mpc_nodes = shared.start_cluster_with_mpc(
17+
2,
18+
2,
19+
1,
20+
contracts.load_mpc_contract(),
21+
)
22+
cluster.init_cluster(mpc_nodes, 2)
23+
cluster.wait_for_state(contract_state.ProtocolState.RUNNING)
24+
25+
yield cluster
26+
27+
atexit._run_exitfuncs()
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Starts 2 near validators and 2 mpc nodes.
4+
Deploys the mpc contract.
5+
Sends ckd requests.
6+
Verifies that the mpc nodes index the ckd request.
7+
Waits for ckd responses. Fails if timeout is reached.
8+
Verifies that ckd responses are correct
9+
"""
10+
11+
import sys
12+
import pathlib
13+
import pytest
14+
15+
sys.path.append(str(pathlib.Path(__file__).resolve().parents[2]))
16+
17+
from common_lib import shared, ckd
18+
from common_lib.constants import CKD_DEPOSIT
19+
20+
21+
@pytest.mark.no_atexit_cleanup
22+
def test_ckd_request_lifecycle(shared_cluster: shared.MpcCluster):
23+
domains = shared_cluster.contract_state().get_running_domains()
24+
25+
bls_domain = None
26+
for domain in domains:
27+
if domain.scheme == "Bls12381":
28+
bls_domain = domain
29+
break
30+
assert bls_domain is not None
31+
32+
keyset = shared_cluster.contract_state().keyset()
33+
assert keyset is not None
34+
public_key = keyset.get_key(bls_domain.id).key["Bls12381"]["public_key"]
35+
36+
app_public_key, app_private_key = ckd.generate_app_public_key()
37+
ckd_args = ckd.generate_ckd_args(bls_domain, app_public_key)
38+
tx = shared_cluster.request_node.sign_tx(
39+
shared_cluster.mpc_contract_account(),
40+
"request_app_private_key",
41+
ckd_args,
42+
deposit=CKD_DEPOSIT,
43+
)
44+
account_id = shared_cluster.request_node.account_id()
45+
46+
tx_hash = shared_cluster.request_node.send_tx(tx)["result"]
47+
res = shared_cluster.request_node.get_tx(tx_hash)
48+
ck = ckd.assert_ckd_success(res)
49+
big_y, big_c = ck["big_y"], ck["big_c"]
50+
51+
assert ckd.verify_ckd(
52+
account_id.encode(), public_key, app_private_key, big_y, big_c
53+
)
Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
#! /usr/bin/env python3
2+
"""
3+
Starts 2 near validators and 2 mpc nodes.
4+
Deploys mpc contract.
5+
Deploys a test contract with a function that makes parallel sign calls.
6+
Calls the test function and ensures a successful response.
7+
"""
8+
9+
import sys
10+
import base64
11+
import pytest
12+
import pathlib
13+
import time
14+
from utils import load_binary_file
15+
from dataclasses import dataclass
16+
17+
sys.path.append(str(pathlib.Path(__file__).resolve().parents[1]))
18+
from common_lib import shared, contracts, constants
19+
from common_lib.shared import metrics
20+
21+
22+
@dataclass
23+
class NodeMetrics:
24+
queue_size: int
25+
requests_indexed: int
26+
responses_indexed: int
27+
matching_responses_indexed: int
28+
29+
def __sub__(self, other):
30+
if isinstance(other, NodeMetrics):
31+
res = NodeMetrics(0, 0, 0, 0)
32+
res.queue_size = self.queue_size - other.queue_size
33+
res.requests_indexed = self.requests_indexed - other.requests_indexed
34+
res.responses_indexed = self.responses_indexed - other.responses_indexed
35+
res.matching_responses_indexed = (
36+
self.matching_responses_indexed - other.matching_responses_indexed
37+
)
38+
return res
39+
40+
41+
def load_parallel_sign_contract() -> bytearray:
42+
"""
43+
Returns test contract for parallel sign
44+
"""
45+
return load_binary_file(contracts.PARALLEL_CONTRACT_BINARY_PATH)
46+
47+
48+
def get_metric_value_for_node(cluster, metric_name: str, node_id: int):
49+
result = cluster.get_int_metric_value_for_node(metric_name, node_id)
50+
return result if result is not None else 0
51+
52+
53+
@pytest.mark.parametrize("num_parallel_requests", [6])
54+
@pytest.mark.no_atexit_cleanup
55+
def test_parallel_sign_calls(
56+
compile_parallel_contract, num_parallel_requests, shared_cluster: shared.MpcCluster
57+
):
58+
assert num_parallel_requests % 3 == 0, "expected number multiple of 3"
59+
# start cluster and deploy mpc contract
60+
contract = load_parallel_sign_contract()
61+
62+
print("Deploying parallel contract")
63+
shared_cluster.deploy_secondary_contract(contract)
64+
65+
started = time.time()
66+
while True:
67+
assert time.time() - started < constants.SHORT_TIMEOUT, "Waiting for metrics"
68+
initial_node_metrics = get_node_metrics_all_nodes(shared_cluster)
69+
initial_queue_attempts = get_queue_attemps_generated(shared_cluster)
70+
if sum(node_metric.queue_size for node_metric in initial_node_metrics) == 0:
71+
break
72+
time.sleep(1)
73+
74+
print("Making parallel request calls")
75+
# call `parallel_sign` and verify that it returns successfully
76+
res = shared_cluster.make_function_call_on_secondary_contract(
77+
function_name="make_parallel_sign_calls",
78+
args={
79+
"target_contract": shared_cluster.mpc_contract_account(),
80+
"ecdsa_calls_by_domain": {0: num_parallel_requests // 3},
81+
"eddsa_calls_by_domain": {1: num_parallel_requests // 3},
82+
"ckd_calls_by_domain": {2: num_parallel_requests // 3},
83+
"seed": 23,
84+
},
85+
)
86+
87+
# check the return value
88+
assert (
89+
"result" in res
90+
and "status" in res["result"]
91+
and "SuccessValue" in res["result"]["status"]
92+
), res
93+
encoded_value = res["result"]["status"]["SuccessValue"]
94+
decoded_value = base64.b64decode(encoded_value).decode("utf-8")
95+
assert int(decoded_value) == num_parallel_requests
96+
97+
target_metrics = NodeMetrics(0, *[num_parallel_requests] * 3)
98+
# check metrics to make sure signature requests are handled properly.
99+
started = time.time()
100+
while True:
101+
assert time.time() - started < constants.SHORT_TIMEOUT, "Waiting for metrics"
102+
metrics_good = True
103+
current_metrics = get_node_metrics_all_nodes(shared_cluster)
104+
for i in range(len(shared_cluster.mpc_nodes)):
105+
if current_metrics[i] - initial_node_metrics[i] != target_metrics:
106+
metrics_good = False
107+
led_requests = (
108+
get_queue_attemps_generated(shared_cluster) - initial_queue_attempts
109+
)
110+
111+
print(f"led_signatures={led_requests}")
112+
if led_requests != num_parallel_requests:
113+
metrics_good = False
114+
if metrics_good:
115+
break
116+
time.sleep(1)
117+
print(
118+
"All requests and responses indexed, all requests had exactly one leader, and signature/ckd queue is empty on all nodes. All Done."
119+
)
120+
121+
122+
def get_node_metrics_all_nodes(cluster: shared.MpcCluster):
123+
number_nodes = len(cluster.mpc_nodes)
124+
125+
network_metrics = [NodeMetrics(0, 0, 0, 0) for _ in range(number_nodes)]
126+
for i in range(len(cluster.mpc_nodes)):
127+
network_metrics[i].queue_size = get_metric_value_for_node(
128+
cluster, "mpc_pending_signatures_queue_size", i
129+
)
130+
network_metrics[i].requests_indexed = get_metric_value_for_node(
131+
cluster, "mpc_pending_signatures_queue_requests_indexed", i
132+
)
133+
network_metrics[i].responses_indexed = get_metric_value_for_node(
134+
cluster, "mpc_pending_signatures_queue_responses_indexed", i
135+
)
136+
network_metrics[i].matching_responses_indexed = get_metric_value_for_node(
137+
cluster, "mpc_pending_signatures_queue_matching_responses_indexed", i
138+
)
139+
140+
network_metrics[i].queue_size += get_metric_value_for_node(
141+
cluster, "mpc_pending_ckds_queue_size", i
142+
)
143+
network_metrics[i].requests_indexed += get_metric_value_for_node(
144+
cluster, "mpc_pending_ckds_queue_requests_indexed", i
145+
)
146+
network_metrics[i].responses_indexed += get_metric_value_for_node(
147+
cluster, "mpc_pending_ckds_queue_responses_indexed", i
148+
)
149+
network_metrics[i].matching_responses_indexed += get_metric_value_for_node(
150+
cluster, "mpc_pending_ckds_queue_matching_responses_indexed", i
151+
)
152+
print(
153+
f"Node {i}: queue_size={network_metrics[i].queue_size}, requests_indexed={network_metrics[i].requests_indexed}, responses_indexed={network_metrics[i].responses_indexed}, matching_responses_indexed={network_metrics[i].matching_responses_indexed}"
154+
)
155+
return network_metrics
156+
157+
158+
def get_queue_attemps_generated(cluster: shared.MpcCluster):
159+
led_requests = cluster.get_int_metric_value(
160+
metrics.IntMetricName.MPC_PENDING_SIGNATURES_QUEUE_ATTEMPTS_GENERATED
161+
) + cluster.get_int_metric_value(
162+
metrics.IntMetricName.MPC_PENDING_CKDS_QUEUE_ATTEMPTS_GENERATED
163+
)
164+
return sum(a for a in led_requests if a is not None)

pytest/tests/test_requests.py renamed to pytest/tests/shared_cluster_tests/test_requests.py

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,35 +15,30 @@
1515

1616
sys.path.append(str(pathlib.Path(__file__).resolve().parents[1]))
1717
from common_lib import shared
18-
from common_lib.contracts import load_mpc_contract
1918
from common_lib.constants import TIMEOUT
2019
from common_lib.shared import metrics
2120

2221

23-
@pytest.mark.parametrize("num_requests, num_respond_access_keys", [(10, 1)])
24-
def test_request_lifecycle(num_requests, num_respond_access_keys):
25-
cluster, mpc_nodes = shared.start_cluster_with_mpc(
26-
2, 2, num_respond_access_keys, load_mpc_contract()
27-
)
28-
cluster.init_cluster(mpc_nodes, 2)
29-
22+
@pytest.mark.parametrize("num_requests", [10])
23+
@pytest.mark.no_atexit_cleanup
24+
def test_request_lifecycle(num_requests, shared_cluster: shared.MpcCluster):
3025
started = time.time()
3126
while True:
3227
time.sleep(1.0)
3328
assert time.time() - started < TIMEOUT, "Waiting for account balances"
3429
# check that the near balance metric works
35-
responder_balances = cluster.get_float_metric_value(
30+
responder_balances = shared_cluster.get_float_metric_value(
3631
metrics.FloatMetricName.MPC_NEAR_RESPONDER_BALANCE
3732
)
3833
print(f"responder_balances: {responder_balances}")
3934
if not all([rb and rb > 0 for rb in responder_balances]):
4035
continue
41-
signer_balances = cluster.get_float_metric_value(
36+
signer_balances = shared_cluster.get_float_metric_value(
4237
metrics.FloatMetricName.MPC_NEAR_SIGNER_BALANCE
4338
)
4439
print(f"signer_balances: {signer_balances}")
4540
if not all([sb and sb > 0 for sb in signer_balances]):
4641
continue
4742
break
48-
cluster.send_and_await_signature_requests(num_requests)
49-
cluster.send_and_await_ckd_requests(num_requests)
43+
shared_cluster.send_and_await_signature_requests(num_requests)
44+
shared_cluster.send_and_await_ckd_requests(num_requests)

0 commit comments

Comments
 (0)