|
1 | 1 | """ |
2 | | -This module contains 'PocketIC', which is the main interface we expose to a test author. |
3 | | -It also contains 'SubnetConfig' and 'SubnetKind', which are used to configure the |
4 | | -subnets of a PocketIC instance. |
| 2 | +This module contains `PocketIC`, which is the main interface exposed to the test author. |
5 | 3 | """ |
6 | 4 |
|
7 | 5 | import base64 |
8 | 6 | import ic |
9 | | -from enum import Enum |
10 | 7 | from ic.candid import Types |
11 | 8 | from typing import Optional, Any |
12 | 9 | from pocket_ic.pocket_ic_server import PocketICServer |
13 | | - |
14 | | - |
15 | | -class SubnetKind(Enum): |
16 | | - """The kind of subnet.""" |
17 | | - |
18 | | - APPLICATION = "Application" |
19 | | - BITCOIN = "Bitcoin" |
20 | | - FIDUCIARY = "Fiduciary" |
21 | | - II = "II" |
22 | | - NNS = "NNS" |
23 | | - SNS = "SNS" |
24 | | - SYSTEM = "System" |
25 | | - VERIFIED_APPLICATION = "VerifiedApplication" |
26 | | - |
27 | | - |
28 | | -class SubnetConfig: |
29 | | - """The configuration of subnets for a PocketIC instance.""" |
30 | | - |
31 | | - def __init__( |
32 | | - self, |
33 | | - application=0, |
34 | | - bitcoin=False, |
35 | | - fiduciary=False, |
36 | | - ii=False, |
37 | | - nns=False, |
38 | | - sns=False, |
39 | | - system=0, |
40 | | - verified_application=0, |
41 | | - ) -> None: |
42 | | - new = {"state_config": "New", "instruction_config": "Production"} |
43 | | - self.application = [new] * application |
44 | | - self.bitcoin = new if bitcoin else None |
45 | | - self.fiduciary = new if fiduciary else None |
46 | | - self.ii = new if ii else None |
47 | | - self.nns = new if nns else None |
48 | | - self.sns = new if sns else None |
49 | | - self.system = [new] * system |
50 | | - self.verified_application = [new] * verified_application |
51 | | - |
52 | | - def __repr__(self) -> str: |
53 | | - return f"SubnetConfigSet(application={self.application}, bitcoin={self.bitcoin}, fiduciary={self.fiduciary}, ii={self.ii}, nns={self.nns}, sns={self.sns}, system={self.system}, verified_application={self.verified_application})" |
54 | | - |
55 | | - def validate(self) -> None: |
56 | | - """Validates the subnet configuration. |
57 | | -
|
58 | | - Raises: |
59 | | - ValueError: if no subnet is configured |
60 | | - """ |
61 | | - if not ( |
62 | | - self.bitcoin |
63 | | - or self.fiduciary |
64 | | - or self.ii |
65 | | - or self.nns |
66 | | - or self.sns |
67 | | - or self.system |
68 | | - or self.application |
69 | | - or self.verified_application |
70 | | - ): |
71 | | - raise ValueError("At least one subnet must be configured.") |
72 | | - |
73 | | - def add_subnet_with_state( |
74 | | - self, subnet_type: SubnetKind, state_dir_path: str, subnet_id: ic.Principal |
75 | | - ): |
76 | | - """Add a subnet with state loaded form the given state directory. |
77 | | - Note that the provided path must be accessible for the PocketIC server process. |
78 | | -
|
79 | | - `state_dir` should point to a directory which is expected to have the following structure: |
80 | | -
|
81 | | - state_dir/ |
82 | | - |-- backups |
83 | | - |-- checkpoints |
84 | | - |-- diverged_checkpoints |
85 | | - |-- diverged_state_markers |
86 | | - |-- fs_tmp |
87 | | - |-- page_deltas |
88 | | - |-- states_metadata.pbuf |
89 | | - |-- tip |
90 | | - `-- tmp |
91 | | -
|
92 | | - `subnet_id` should be the subnet ID of the subnet in the state to be loaded""" |
93 | | - |
94 | | - raw_subnet_id_bytes = base64.b64encode( |
95 | | - subnet_id.bytes |
96 | | - ) # convert bytes to base64 bytes |
97 | | - raw_subnet_id = raw_subnet_id_bytes.decode() # convert bytes to str |
98 | | - |
99 | | - new_from_path = { |
100 | | - "state_config": { |
101 | | - "FromPath": [state_dir_path, {"subnet_id": raw_subnet_id}] |
102 | | - }, |
103 | | - "instruction_config": "Production", |
104 | | - } |
105 | | - |
106 | | - match subnet_type: |
107 | | - case SubnetKind.APPLICATION: |
108 | | - self.application.append(new_from_path) |
109 | | - case SubnetKind.BITCOIN: |
110 | | - self.bitcoin = new_from_path |
111 | | - case SubnetKind.FIDUCIARY: |
112 | | - self.fiduciary = new_from_path |
113 | | - case SubnetKind.II: |
114 | | - self.ii = new_from_path |
115 | | - case SubnetKind.NNS: |
116 | | - self.nns = new_from_path |
117 | | - case SubnetKind.SNS: |
118 | | - self.sns = new_from_path |
119 | | - case SubnetKind.SYSTEM: |
120 | | - self.system.append(new_from_path) |
121 | | - case SubnetKind.VERIFIED_APPLICATION: |
122 | | - self.verified_application.append(new_from_path) |
123 | | - |
124 | | - def _json(self) -> dict: |
125 | | - return { |
126 | | - "subnet_config_set": { |
127 | | - "application": self.application, |
128 | | - "bitcoin": self.bitcoin, |
129 | | - "fiduciary": self.fiduciary, |
130 | | - "ii": self.ii, |
131 | | - "nns": self.nns, |
132 | | - "sns": self.sns, |
133 | | - "system": self.system, |
134 | | - "verified_application": self.verified_application, |
135 | | - }, |
136 | | - "state_dir": None, |
137 | | - "nonmainnet_features": False, |
138 | | - "log_level": None, |
139 | | - "bitcoind_addr": None, |
140 | | - } |
| 10 | +from pocket_ic.subnet_config import SubnetConfig, SubnetKind |
141 | 11 |
|
142 | 12 |
|
143 | 13 | class PocketIC: |
@@ -549,19 +419,21 @@ def update_call_with_effective_principal( |
549 | 419 | "payload": base64.b64encode(payload).decode(), |
550 | 420 | } |
551 | 421 |
|
552 | | - submit_ingress_message = self._instance_post("update/submit_ingress_message", body) |
| 422 | + submit_ingress_message = self._instance_post( |
| 423 | + "update/submit_ingress_message", body |
| 424 | + ) |
553 | 425 | ok = self._get_ok(submit_ingress_message) |
554 | 426 | result = self._instance_post("update/await_ingress_message", ok) |
555 | 427 | return self._get_ok_data(result) |
556 | | - |
| 428 | + |
557 | 429 | def _get_ok(self, request_result): |
558 | 430 | if "Ok" in request_result: |
559 | 431 | return request_result["Ok"] |
560 | 432 | if "Err" in request_result: |
561 | | - err = request_result['Err'] |
562 | | - reject_code = err['reject_code'] |
563 | | - reject_message = err['reject_message'] |
564 | | - error_code = err['error_code'] |
| 433 | + err = request_result["Err"] |
| 434 | + reject_code = err["reject_code"] |
| 435 | + reject_message = err["reject_message"] |
| 436 | + error_code = err["error_code"] |
565 | 437 | msg = f"PocketIC returned a rejection error: reject code {reject_code}, reject message {reject_message}, error code {error_code}" |
566 | 438 | raise ValueError(msg) |
567 | 439 | raise ValueError(f"Malformed response: {request_result}") |
|
0 commit comments