From c2e3945f3cb1456767d82f04f431f8814ff27bde Mon Sep 17 00:00:00 2001 From: Fatih Acar Date: Sun, 4 Feb 2024 21:43:37 +0100 Subject: [PATCH] feat(tests, scale): add rels parameter This parameter allows to add relationships to the node used for scale testing. Signed-off-by: Fatih Acar --- backend/tests/scale/common/stagers.py | 35 +++++++++++++++++++++------ backend/tests/scale/main.py | 35 +++++++++++++++++++++------ tasks/backend.py | 6 ++++- 3 files changed, 61 insertions(+), 15 deletions(-) diff --git a/backend/tests/scale/common/stagers.py b/backend/tests/scale/common/stagers.py index 2dadb140b5..c07c55352f 100644 --- a/backend/tests/scale/common/stagers.py +++ b/backend/tests/scale/common/stagers.py @@ -7,21 +7,42 @@ def load_schema( - client: InfrahubClientSync, schema: Path, branch: Optional[str] = None, extra_attributes: List[dict] = [] + client: InfrahubClientSync, + schema: Path, + branch: Optional[str] = None, + extra_attributes: List[dict] = [], + relationships: List[dict] = [], ): branch = branch or "main" data = yaml.safe_load(schema.read_text()) - for attr in extra_attributes: - data["nodes"][0]["attributes"].append({"name": attr["name"], "kind": attr["kind"]}) + + data["nodes"][0]["attributes"] += extra_attributes + + if "relationships" not in data["nodes"][0]: + data["nodes"][0]["relationships"] = list() + data["nodes"][0]["relationships"] += relationships + client.schema.validate(data) - client.schema.load(schemas=[data], branch=branch) + (loaded, response) = client.schema.load(schemas=[data], branch=branch) + if not loaded: + raise ValueError(f"Could not load schema: {response}") -def _stage_node(client: InfrahubClientSync, kind: str, prefix: str, amount: int, attrs: int, offset: int = 0): +def _stage_node( + client: InfrahubClientSync, kind: str, prefix: str, amount: int, attrs: int, rels: int, offset: int = 0 +): client.schema.get("InfraNode") extra_attributes = dict() for i in range(attrs): - extra_attributes[f"test{i}"] = "test data" + extra_attributes[f"attr{i}"] = "test data" + + # Create a tag to use for relationship + if rels > 0: + tag = client.create(kind="BuiltinTag", data={"name": "testtag"}) + tag.save() + for i in range(rels): + extra_attributes[f"rel{i}"] = tag.id + for i in range(offset, offset + amount): node = client.create(kind=kind, data={"name": f"{prefix}{i}", **extra_attributes}) node.save() @@ -30,7 +51,7 @@ def _stage_node(client: InfrahubClientSync, kind: str, prefix: str, amount: int, stage_infranode = partial(_stage_node, kind="InfraNode", prefix="Node") -def _stage_branch(client: InfrahubClientSync, prefix: str, amount: int, attrs: int, offset: int = 0): +def _stage_branch(client: InfrahubClientSync, prefix: str, amount: int, attrs: int, rels: int, offset: int = 0): for i in range(offset, offset + amount): client.branch.create(branch_name=f"{prefix}{i}", description="description", data_only=True) diff --git a/backend/tests/scale/main.py b/backend/tests/scale/main.py index e87e11f4f9..6280c59bbf 100644 --- a/backend/tests/scale/main.py +++ b/backend/tests/scale/main.py @@ -13,7 +13,7 @@ config = ScaleTestConfig() -def stage_environment(function: str, amount: int, attrs: int, schema: Path): +def stage_environment(function: str, amount: int, attrs: int, rels: int, schema: Path): if function == "": return @@ -28,16 +28,31 @@ def stage_environment(function: str, amount: int, attrs: int, schema: Path): address=config.url, config=Config(api_token=config.api_token, timeout=config.client_timeout) ) print("--- loading load testing schema") + + # Generate extra attributes attributes = [] for i in range(attrs): - attributes.append({"name": f"test{i}", "kind": "Text", "default_value": "", "optional": True}) - - common.stagers.load_schema(staging_client, schema, extra_attributes=attributes) + attributes.append({"name": f"attr{i}", "kind": "Text", "default_value": "", "optional": True}) + # Generate extra relationships + relationships = [] + for i in range(rels): + relationships.append( + { + "name": f"rel{i}", + "kind": "Generic", + "peer": "BuiltinTag", + "cardinality": "one", + "identifier": f"builtintag__infranode_{i}", + "optional": True, + } + ) + + common.stagers.load_schema(staging_client, schema, extra_attributes=attributes, relationships=relationships) print("--- done") time.sleep(5) print("--- staging load testing environment") - stager(client=staging_client, amount=amount, attrs=attrs) + stager(client=staging_client, amount=amount, attrs=attrs, rels=rels) print("--- 20s cool down period") time.sleep(20) @@ -63,15 +78,21 @@ def stage_environment(function: str, amount: int, attrs: int, schema: Path): type=click.IntRange(min=0, max=1_000_000_000), help="Amount of attributes per object to be created in the `staging function`", ) +@click.option( + "--rels", + default=0, + type=click.IntRange(min=0, max=1_000_000_000), + help="Amount of relationships per object to be created in the `staging function`", +) @click.option("--test", default="InfrahubClientUser", help="The Locust test user class") -def main(schema: Path, stager: str, amount: int, attrs: int, test: str) -> int: +def main(schema: Path, stager: str, amount: int, attrs: int, rels: int, test: str) -> int: if not hasattr(common.users, test): print(f"Invalid test class provided: {test}") return 1 user_class = getattr(common.users, test) - stage_environment(function=stager, amount=amount, attrs=attrs, schema=schema) + stage_environment(function=stager, amount=amount, attrs=attrs, rels=rels, schema=schema) print("--- starting test") env = Environment(user_classes=[user_class], events=events) diff --git a/tasks/backend.py b/tasks/backend.py index a08f5a60be..8166fbcc01 100644 --- a/tasks/backend.py +++ b/tasks/backend.py @@ -165,7 +165,7 @@ def test_scale_env_destroy(context: Context, database: str = INFRAHUB_DATABASE): return execute_command(context=context, command=command) -@task(optional=["schema", "stager", "amount", "test", "attrs"]) +@task(optional=["schema", "stager", "amount", "test", "attrs", "rels"]) def test_scale( context: Context, schema: Path = f"{ESCAPED_REPO_PATH}/backend/tests/scale/schema.yml", @@ -173,6 +173,7 @@ def test_scale( amount: int = None, test: str = None, attrs: int = None, + rels: int = None, ): args = [] if stager: @@ -190,6 +191,9 @@ def test_scale( if attrs: args.extend(["--attrs", attrs]) + if rels: + args.extend(["--rels", rels]) + with context.cd(ESCAPED_REPO_PATH): base_cmd = ["python", "backend/tests/scale/main.py"] cmd = " ".join(base_cmd + args)