2025-02-14 11:00:17 -03:00

84 lines
2.3 KiB
Python

from io import StringIO
from typing import cast
import yaml
from urllib3.util import parse_url
from benchmarks.codex.codex_node import CodexNode
from benchmarks.codex.config import (
CodexNodeConfig,
CodexNodeSetConfig,
CodexExperimentConfig,
)
def test_should_expand_node_sets_into_simple_nodes():
nodeset = CodexNodeSetConfig(
network_size=3,
first_node_index=0,
name="codex-{node_index}",
address="codex-{node_index}.local.svc",
disc_port=6890,
api_port=6891,
agent_url="http://codex-{node_index}.local.svc:9000",
)
assert nodeset.nodes == [
CodexNodeConfig(
name="codex-0",
address="codex-0.local.svc",
disc_port=6890,
api_port=6891,
agent_url="http://codex-0.local.svc:9000",
),
CodexNodeConfig(
name="codex-1",
address="codex-1.local.svc",
disc_port=6890,
api_port=6891,
agent_url="http://codex-1.local.svc:9000",
),
CodexNodeConfig(
name="codex-2",
address="codex-2.local.svc",
disc_port=6890,
api_port=6891,
agent_url="http://codex-2.local.svc:9000",
),
]
def test_should_build_experiment_from_config():
config_file = StringIO("""
codex_experiment:
repetitions: 3
seeders: 3
seeder_sets: 3
file_size: 1024
logging_cooldown: 10
nodes:
network_size: 5
first_node_index: 0
name: "codex-nodes-{node_index}"
address: "codex-nodes-{node_index}.codex-nodes-service.codex-benchmarks.svc.cluster.local"
disc_port: 6890
api_port: 6891
agent_url: "http://codex-nodes-{node_index}.codex-nodes-service.codex-benchmarks.svc.cluster.local:9000/"
""")
config = CodexExperimentConfig.model_validate(
yaml.safe_load(config_file)["codex_experiment"]
)
experiment = config.build()
repetitions = list(experiment.experiments)
assert len(repetitions) == 9
assert len(repetitions[0].experiment.nodes) == 5
assert cast(
CodexNode, repetitions[0].experiment.nodes[4]
).codex_api_url == parse_url(
"http://codex-nodes-4.codex-nodes-service.codex-benchmarks.svc.cluster.local:6891"
)