feat(tests, scale): add rels parameter

This parameter allows to add relationships to the node used for scale
testing.

Signed-off-by: Fatih Acar <[email protected]>

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)
 

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)

tasks/backend.py

@@ -165,14 +165,15 @@ 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",
     stager: str = None,
     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)