diff --git a/src/benchmark_scenarios/private_chats.py b/src/benchmark_scenarios/private_chats.py index 5fab782..962733b 100644 --- a/src/benchmark_scenarios/private_chats.py +++ b/src/benchmark_scenarios/private_chats.py @@ -8,7 +8,7 @@ import src.logger from src import kube_utils, setup_status from src.inject_messages import inject_messages_one_to_one from src.setup_status import initialize_nodes_application, send_friend_requests, accept_friend_requests, \ - decline_friend_requests + decline_friend_requests, create_group_chat, add_contacts logger = logging.getLogger(__name__) @@ -129,3 +129,32 @@ async def send_one_to_one_message(): logger.info("Shutting down node connections") await asyncio.gather(*[node.shutdown() for node in relay_nodes.values()]) logger.info("Finished send_one_to_one_message") + + +async def create_private_group(): + # 10 admin nodes + # 100 single-group members + # Each admin node create a group and invite 10 single-group members in it, who accept the invite + # single-group members accept the invite + kube_utils.setup_kubernetes_client() + backend_relay_pods = kube_utils.get_pods("status-backend-relay", "status-go-test") + relay_nodes = await initialize_nodes_application(backend_relay_pods) + + backend_relay_pods = [pod_name.split(".")[0] for pod_name in backend_relay_pods] + + admin_nodes = backend_relay_pods[:10] + members = backend_relay_pods[10:110] + members_pub_keys = [relay_nodes[node].public_key for node in members] + + # In order to create a group, first they need to be friends + friend_requests = await send_friend_requests(relay_nodes, admin_nodes, members) + logger.info("Accepting friend requests") + _ = await accept_friend_requests(relay_nodes, friend_requests) + _ = await add_contacts(relay_nodes, admin_nodes, members) + + # 1 admin to 10 users, no overlap + await asyncio.gather(*[create_group_chat(relay_nodes[admin], members_pub_keys[10*i: (10*i)+10]) for i, admin in enumerate(admin_nodes)]) + + logger.info("Shutting down node connections") + await asyncio.gather(*[node.shutdown() for node in relay_nodes.values()]) + logger.info("Finished create_private_group") diff --git a/src/setup_status.py b/src/setup_status.py index e40b490..25628a2 100644 --- a/src/setup_status.py +++ b/src/setup_status.py @@ -1,6 +1,8 @@ # Python Imports import asyncio import logging +import random +import string import time # Project Imports @@ -171,6 +173,16 @@ async def accept_friend_requests(nodes: dict[str, StatusBackend], requests: list logger.info(f"All {total_requests} friend requests accepted.") +async def add_contacts(nodes: dict[str, StatusBackend], adders: list[str], contacts: list[str]): + async def _add_contacts_to_node(nodes: dict[str, StatusBackend], adder: str, contacts: list[str]): + _ = await asyncio.gather(*[nodes[adder].wakuext_service.add_contact(nodes[contact].public_key, contact) for contact in contacts]) + + return _ + + _ = await asyncio.gather(*[_add_contacts_to_node(nodes, adder, contacts) for adder in adders]) + + logger.info(f"All {len(contacts)} contacts added to {len(adders)} nodes.") + async def decline_friend_requests(nodes: dict[str, StatusBackend], requests: list[(str, dict[str, str])]): # Flatten all tasks into a single list and execute them concurrently @@ -202,6 +214,15 @@ async def decline_friend_requests(nodes: dict[str, StatusBackend], requests: lis logger.info(f"All {total_requests} friend requests rejected.") +async def create_group_chat(admin: StatusBackend, receivers: list[str]): + name = f"private_group_{''.join(random.choices(string.ascii_letters, k=10))}" + logger.info(f"Creating private group {name}") + response = await admin.wakuext_service.create_group_chat_with_members(receivers, name) + group_id = response.get("result", {}).get("communities", [{}])[0].get("id") + logger.info(f"Group {name} created with ID {group_id}") + + + async def get_messages_by_content_type(response: dict, content_type: str, message_pattern: str="") -> list[dict]: matched_messages = [] messages = response.get("result", {}).get("messages", [])