Skip to content

Commit f2a3222

Browse files
committed
scenarios: use initContainer
1 parent 2ab431b commit f2a3222

File tree

10 files changed

+87
-26
lines changed

10 files changed

+87
-26
lines changed

resources/charts/commander/templates/pod.yaml

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,30 @@ metadata:
88
mission: commander
99
spec:
1010
restartPolicy: {{ .Values.restartPolicy }}
11+
initContainers:
12+
- name: init
13+
image: busybox
14+
command: ["/bin/sh", "-c"]
15+
args:
16+
- |
17+
while [ ! -f /shared/archive.pyz ]; do
18+
echo "Waiting for /shared/archive.pyz to exist..."
19+
sleep 2
20+
done
21+
volumeMounts:
22+
- name: shared-volume
23+
mountPath: /shared
1124
containers:
1225
- name: {{ .Chart.Name }}
1326
image: python:3.12-slim
1427
imagePullPolicy: IfNotPresent
1528
command: ["/bin/sh", "-c"]
1629
args:
1730
- |
18-
python3 /archive.pyz {{ .Values.args }}
31+
python3 /shared/archive.pyz {{ .Values.args }}
1932
volumeMounts:
20-
- name: warnet
21-
mountPath: /warnet.json
22-
subPath: warnet.json
33+
- name: shared-volume
34+
mountPath: /shared
2335
volumes:
24-
- name: warnet
25-
configMap:
26-
name: {{ include "commander.fullname" . }}-warnet
36+
- name: shared-volume
37+
emptyDir: {}

resources/scenarios/commander.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import signal
99
import sys
1010
import tempfile
11-
from pathlib import Path
1211
from typing import Dict
1312

1413
from test_framework.authproxy import AuthServiceProxy
@@ -21,7 +20,7 @@
2120
from test_framework.test_node import TestNode
2221
from test_framework.util import PortSeed, get_rpc_proxy
2322

24-
WARNET_FILE = Path("/warnet.json")
23+
WARNET_FILE = "/shared/warnet.json"
2524

2625
try:
2726
with open(WARNET_FILE) as file:

src/warnet/control.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import base64
21
import io
32
import json
43
import os
@@ -26,6 +25,8 @@
2625
pod_log,
2726
snapshot_bitcoin_datadir,
2827
wait_for_pod,
28+
wait_for_init,
29+
write_file_to_container,
2930
)
3031
from .process import run_command, stream_command
3132

@@ -202,7 +203,7 @@ def run(scenario_file: str, debug: bool, additional_args: tuple[str]):
202203
]
203204

204205
# Encode tank data for warnet.json
205-
warnet_data = base64.b64encode(json.dumps(tanks).encode()).decode()
206+
warnet_data = json.dumps(tanks).encode()
206207

207208
# Create in-memory buffer to store python archive instead of writing to disk
208209
archive_buffer = io.BytesIO()
@@ -226,8 +227,9 @@ def filter(path):
226227

227228
# Encode the binary data as Base64
228229
archive_buffer.seek(0)
229-
archive_data = base64.b64encode(archive_buffer.read()).decode()
230+
archive_data = archive_buffer.read()
230231

232+
# Start the commander pod with python and init containers
231233
try:
232234
# Construct Helm command
233235
helm_command = [
@@ -238,8 +240,6 @@ def filter(path):
238240
namespace,
239241
"--set",
240242
f"fullnameOverride={name}",
241-
"--set",
242-
f"warnet={warnet_data}",
243243
]
244244

245245
# Add additional arguments
@@ -252,19 +252,24 @@ def filter(path):
252252
result = subprocess.run(helm_command, check=True, capture_output=True, text=True)
253253

254254
if result.returncode == 0:
255-
print(f"Successfully started scenario: {scenario_name}")
255+
print(f"Successfully deployed scenario commander: {scenario_name}")
256256
print(f"Commander pod name: {name}")
257257
else:
258-
print(f"Failed to start scenario: {scenario_name}")
258+
print(f"Failed to deploy scenario commander: {scenario_name}")
259259
print(f"Error: {result.stderr}")
260260

261261
except subprocess.CalledProcessError as e:
262-
print(f"Failed to start scenario: {scenario_name}")
262+
print(f"Failed to deploy scenario commander: {scenario_name}")
263263
print(f"Error: {e.stderr}")
264264

265+
# upload scenario files and network data to the init container
266+
wait_for_init(name)
267+
if write_file_to_container(
268+
name, "init", "/shared/warnet.json", warnet_data
269+
) and write_file_to_container(name, "init", "/shared/archive.pyz", archive_data):
270+
print(f"Successfully uploaded scenario data to commander: {scenario_name}")
271+
265272
if debug:
266-
print("Waiting for commander pod to start...")
267-
wait_for_pod(name)
268273
_logs(pod_name=name, follow=True)
269274
print("Deleting pod...")
270275
delete_pod(name)

src/warnet/k8s.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,24 @@ def wait_for_pod_ready(name, namespace, timeout=300):
264264
return False
265265

266266

267+
def wait_for_init(pod_name, timeout=300):
268+
sclient = get_static_client()
269+
namespace = get_default_namespace()
270+
w = watch.Watch()
271+
for event in w.stream(
272+
sclient.list_namespaced_pod, namespace=namespace, timeout_seconds=timeout
273+
):
274+
pod = event["object"]
275+
if pod.metadata.name == pod_name:
276+
for init_container_status in pod.status.init_container_statuses:
277+
if init_container_status.state.running:
278+
print(f"initContainer in pod {pod_name} is ready")
279+
w.stop()
280+
return True
281+
print(f"Timeout waiting for initContainer in {pod_name} to be ready.")
282+
return False
283+
284+
267285
def wait_for_ingress_controller(timeout=300):
268286
# get name of ingress controller pod
269287
sclient = get_static_client()
@@ -308,3 +326,28 @@ def wait_for_pod(pod_name, timeout_seconds=10):
308326
return
309327
sleep(1)
310328
timeout_seconds -= 1
329+
330+
331+
def write_file_to_container(pod_name, container_name, dst_path, data):
332+
sclient = get_static_client()
333+
namespace = get_default_namespace()
334+
exec_command = ["sh", "-c", f"cat > {dst_path}"]
335+
try:
336+
res = stream(
337+
sclient.connect_get_namespaced_pod_exec,
338+
pod_name,
339+
namespace,
340+
command=exec_command,
341+
container=container_name,
342+
stdin=True,
343+
stderr=True,
344+
stdout=True,
345+
tty=False,
346+
_preload_content=False,
347+
)
348+
res.write_stdin(data)
349+
res.close()
350+
print(f"Successfully copied data to {pod_name}({container_name}):{dst_path}")
351+
return True
352+
except Exception as e:
353+
print(f"Failed to copy data to {pod_name}({container_name}):{dst_path}:\n{e}")

src/warnet/network.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ def copy_scenario_defaults(directory: Path):
4444
directory,
4545
SCENARIOS_DIR.name,
4646
SCENARIOS_DIR,
47-
["__pycache__", "TEST_*.py"],
47+
["__pycache__", "testscenario_*.py"],
4848
)
4949

5050

test/dag_connection_test.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ class DAGConnectionTest(TestBase):
1010
def __init__(self):
1111
super().__init__()
1212
self.network_dir = Path(os.path.dirname(__file__)) / "data" / "ten_semi_unconnected"
13+
self.scen_dir = Path(os.path.dirname(__file__)).parent / "resources" / "scenarios"
1314

1415
def run_test(self):
1516
try:
@@ -25,8 +26,9 @@ def setup_network(self):
2526
self.wait_for_all_edges()
2627

2728
def run_connect_dag_scenario(self):
28-
self.log.info("Running connect_dag scenario")
29-
self.warnet("run resources/scenarios/TEST_connect_dag.py")
29+
scenario_file = self.scen_dir / "testscenario_connect_dag.py"
30+
self.log.info(f"Running scenario from: {scenario_file}")
31+
self.warnet(f"run {scenario_file}")
3032
self.wait_for_all_scenarios()
3133

3234

test/scenarios_test.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ class ScenariosTest(TestBase):
1414
def __init__(self):
1515
super().__init__()
1616
self.network_dir = Path(os.path.dirname(__file__)) / "data" / "12_node_ring"
17+
self.scen_dir = Path(os.path.dirname(__file__)).parent / "resources" / "scenarios"
1718

1819
def run_test(self):
1920
try:
@@ -71,7 +72,7 @@ def check_blocks(self, target_blocks, start: int = 0):
7172
return count >= start + target_blocks
7273

7374
def run_and_check_miner_scenario_from_file(self):
74-
scenario_file = "resources/scenarios/miner_std.py"
75+
scenario_file = self.scen_dir / "miner_std.py"
7576
self.log.info(f"Running scenario from file: {scenario_file}")
7677
self.warnet(f"run {scenario_file} --allnodes --interval=1")
7778
start = int(self.warnet("bitcoin rpc tank-0000 getblockcount"))
@@ -82,19 +83,19 @@ def run_and_check_miner_scenario_from_file(self):
8283
self.stop_scenario()
8384

8485
def run_and_check_scenario_from_file(self):
85-
scenario_file = "resources/scenarios/TEST_p2p_interface.py"
86+
scenario_file = self.scen_dir / "testscenario_p2p_interface.py"
8687
self.log.info(f"Running scenario from: {scenario_file}")
8788
self.warnet(f"run {scenario_file}")
8889
self.wait_for_predicate(self.check_scenario_clean_exit)
8990

9091
def check_regtest_recon(self):
91-
scenario_file = "resources/scenarios/reconnaissance.py"
92+
scenario_file = self.scen_dir / "reconnaissance.py"
9293
self.log.info(f"Running scenario from file: {scenario_file}")
9394
self.warnet(f"run {scenario_file}")
9495
self.wait_for_predicate(self.check_scenario_clean_exit)
9596

9697
def check_active_count(self):
97-
scenario_file = "resources/scenarios/TEST_buggy_failure.py"
98+
scenario_file = self.scen_dir / "testscenario_buggy_failure.py"
9899
self.log.info(f"Running scenario from: {scenario_file}")
99100
self.warnet(f"run {scenario_file}")
100101

0 commit comments

Comments
 (0)