Skip to content

Commit 46648c2

Browse files
authored
feat: Update PocketIC to support latest server, load/store state dirs (#71)
- factor out SubnetConfig into separate file - add functionality to load/store state to state dir - support PocketIC server version 9.0
1 parent fd6beb6 commit 46648c2

File tree

8 files changed

+224
-158
lines changed

8 files changed

+224
-158
lines changed

.github/workflows/build.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ jobs:
2323
- name: Install PocketIC server
2424
uses: dfinity/pocketic@main
2525
with:
26-
pocket-ic-server-version: "8.0.0"
26+
pocket-ic-server-version: "9.0.0"
2727

2828
- name: Create python environment and run tests and examples
2929
run: |

CHANGELOG.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
88

99
## Unreleased
1010

11+
## 3.1.0 - 2025-04-28
12+
13+
### Added
14+
- Support for PocketIC server version 9.0.0
15+
- Load and store PocketIC state from/to a directory
16+
17+
18+
1119
## 3.0.1 - 2025-02-27
1220

1321
### Added

pocket_ic/__init__.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
"""
2-
This package contains 'PocketIC' and 'PocketICServer'.
2+
`PocketIC` represents an IC instance and is the public interface
3+
of this package.
34
4-
'PocketIC' represents an IC instance and is the public interface
5-
of this package.
5+
Instances are running on a PocketIC server, which is represented
6+
by `PocketICServer`.
67
7-
Instances are running on a PocketIC server, which is represented
8-
by 'PocketICServer'.
8+
`SubnetConfig` is used to configure the subnets of a PocketIC instance.
99
"""
10+
1011
from .pocket_ic import *
1112
from .pocket_ic_server import *
13+
from .subnet_config import *

pocket_ic/pocket_ic.py

Lines changed: 10 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -1,143 +1,13 @@
11
"""
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.
53
"""
64

75
import base64
86
import ic
9-
from enum import Enum
107
from ic.candid import Types
118
from typing import Optional, Any
129
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
14111

14212

14313
class PocketIC:
@@ -549,19 +419,21 @@ def update_call_with_effective_principal(
549419
"payload": base64.b64encode(payload).decode(),
550420
}
551421

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+
)
553425
ok = self._get_ok(submit_ingress_message)
554426
result = self._instance_post("update/await_ingress_message", ok)
555427
return self._get_ok_data(result)
556-
428+
557429
def _get_ok(self, request_result):
558430
if "Ok" in request_result:
559431
return request_result["Ok"]
560432
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"]
565437
msg = f"PocketIC returned a rejection error: reject code {reject_code}, reject message {reject_message}, error code {error_code}"
566438
raise ValueError(msg)
567439
raise ValueError(f"Malformed response: {request_result}")

pocket_ic/pocket_ic_server.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
"""
2-
This module contains the 'PocketICServer', which starts or discovers a PocketIC server process.
2+
This module contains the `PocketICServer`, which starts or discovers a PocketIC server process.
33
"""
44

55
import os
66
import time
77
import requests
8-
from typing import List, Tuple, Optional
8+
from typing import List, Optional
99
from tempfile import gettempdir
1010

1111

pocket_ic/subnet_config.py

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
"""
2+
This module contains `SubnetConfig` and `SubnetKind`, which are used to configure the
3+
subnets of a PocketIC instance.
4+
"""
5+
6+
from enum import Enum
7+
import os
8+
9+
10+
class SubnetKind(Enum):
11+
"""The kind of subnet."""
12+
13+
APPLICATION = "Application"
14+
BITCOIN = "Bitcoin"
15+
FIDUCIARY = "Fiduciary"
16+
II = "II"
17+
NNS = "NNS"
18+
SNS = "SNS"
19+
SYSTEM = "System"
20+
VERIFIED_APPLICATION = "VerifiedApplication"
21+
22+
23+
class SubnetConfig:
24+
"""The configuration of subnets for a PocketIC instance.
25+
26+
`state_dir`:
27+
Use `state_dir=<empty_dir>` to store the state of the subnets in a given directory.
28+
This directory can be used later with `state_dir=<dir>` to load a previous state.
29+
Note that the provided path must be accessible for the PocketIC server process.
30+
"""
31+
32+
def __init__(
33+
self,
34+
application=0,
35+
bitcoin=False,
36+
fiduciary=False,
37+
ii=False,
38+
nns=False,
39+
sns=False,
40+
system=0,
41+
verified_application=0,
42+
state_dir: str | None = None,
43+
) -> None:
44+
new = {"state_config": "New", "instruction_config": "Production"}
45+
self.application = [new] * application
46+
self.bitcoin = new if bitcoin else None
47+
self.fiduciary = new if fiduciary else None
48+
self.ii = new if ii else None
49+
self.nns = new if nns else None
50+
self.sns = new if sns else None
51+
self.system = [new] * system
52+
self.verified_application = [new] * verified_application
53+
self.state_dir = state_dir
54+
55+
def __repr__(self) -> str:
56+
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})"
57+
58+
def validate(self) -> None:
59+
"""Validates the subnet configuration.
60+
61+
Raises:
62+
ValueError: if no subnet and no state dir is configured
63+
"""
64+
if not (
65+
self.bitcoin
66+
or self.fiduciary
67+
or self.ii
68+
or self.nns
69+
or self.sns
70+
or self.system
71+
or self.application
72+
or self.verified_application
73+
or len(os.listdir(self.state_dir)) != 0
74+
):
75+
raise ValueError(
76+
"At least one subnet or a non-empty state directory must be configured."
77+
)
78+
79+
def add_subnet_with_state(self, subnet_type: SubnetKind, state_dir_path: str):
80+
"""Add a single subnet with state loaded form the given state directory.
81+
Note that the provided path must be accessible for the PocketIC server process.
82+
83+
`state_dir` should point to a directory which is expected to have the following structure:
84+
85+
state_dir/
86+
|-- backups
87+
|-- checkpoints
88+
|-- diverged_checkpoints
89+
|-- diverged_state_markers
90+
|-- fs_tmp
91+
|-- page_deltas
92+
|-- states_metadata.pbuf
93+
|-- tip
94+
|-- tmp
95+
"""
96+
97+
new_from_path = {
98+
"state_config": {"FromPath": state_dir_path},
99+
"instruction_config": "Production",
100+
}
101+
102+
match subnet_type:
103+
case SubnetKind.APPLICATION:
104+
self.application.append(new_from_path)
105+
case SubnetKind.BITCOIN:
106+
self.bitcoin = new_from_path
107+
case SubnetKind.FIDUCIARY:
108+
self.fiduciary = new_from_path
109+
case SubnetKind.II:
110+
self.ii = new_from_path
111+
case SubnetKind.NNS:
112+
self.nns = new_from_path
113+
case SubnetKind.SNS:
114+
self.sns = new_from_path
115+
case SubnetKind.SYSTEM:
116+
self.system.append(new_from_path)
117+
case SubnetKind.VERIFIED_APPLICATION:
118+
self.verified_application.append(new_from_path)
119+
120+
def _json(self) -> dict:
121+
return {
122+
"subnet_config_set": {
123+
"application": self.application,
124+
"bitcoin": self.bitcoin,
125+
"fiduciary": self.fiduciary,
126+
"ii": self.ii,
127+
"nns": self.nns,
128+
"sns": self.sns,
129+
"system": self.system,
130+
"verified_application": self.verified_application,
131+
},
132+
"state_dir": self.state_dir,
133+
"nonmainnet_features": False,
134+
"log_level": None,
135+
"bitcoind_addr": None,
136+
}

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "pocket_ic"
3-
version = "3.0.1"
3+
version = "3.1.0"
44
description = "PocketIC: A Canister Smart Contract Testing Platform"
55
authors = [
66
"The Internet Computer Project Developers <dept-testing_&_verification@dfinity.org>",

0 commit comments

Comments
 (0)