diff --git a/utils/containers.py b/utils/containers.py index 5737800419df035d9aad76973da9f7302f7e3833..e923e803bb20770aa01155e3e4f4d13f7a1ae396 100644 --- a/utils/containers.py +++ b/utils/containers.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, List from dataclasses import dataclass import os import json @@ -50,6 +50,9 @@ class ClusterNode: def external_endpoint(self) -> str: return f"http://localhost:{self.external_port}" + def node_info(self) -> dict: + return {"address": f"{self.ip}:{self.port}", "id": self.index} + class ClusterConductor: # _parent: KVSTestFixture @@ -115,8 +118,7 @@ class ClusterConductor: for container in containers: if container and container_regex.match(container): self._dump_container_logs(dir, container) - def get_view(self) -> str: - return {"address": f"{self.ip}:{self.port}", "id": self.index} + def _dump_container_logs(self, dir, name: str) -> None: log_file = os.path.join(dir, f"{name}.log") self.log(f"Dumping logs for container {name} to file {log_file}") @@ -507,10 +509,7 @@ class ClusterConductor: def get_node(self, index): return self.nodes[index] def get_full_view(self): - view = [] - for node in self.nodes: - view.append({"address": f"{node.ip}:{node.port}", "id": node.index}) - return view + return self.get_shard_view() def get_node(self, index): return self.nodes[index] @@ -553,19 +552,18 @@ class ClusterConductor: return self.shards[shard_name].remove(node) - def get_shard_view(self) -> dict: + def get_shard_view(self) -> dict[str, list[dict[str, Any]]]: return { - shard: [node.get_view() for node in nodes] + shard: [node.node_info() for node in nodes] for shard, nodes in self.shards.items() } def get_partition_view(self, partition_id: str): net_name = f"kvs_{self.group_id}_net_{partition_id}" - view = [] - for node in self.nodes: - if net_name in node.networks: - view.append({"address": f"{node.ip}:{node.port}", "id": node.index}) - return view + return { + shard: [node.node_info() for node in nodes if net_name in node.networks] + for shard, nodes in self.shards.items() + } def __enter__(self): return self