feat: add load test

loadTests/loadtest.js

@@ -0,0 +1,213 @@
+import http from "k6/http";
+import { sleep, check } from "k6";
+
+let target = __ENV.TARGET ? parseInt(__ENV.TARGET, 10) : 50;
+let host = __ENV.HOST || "http://localhost:8081";
+let action = __ENV.ACTION || "AGGREGATED";
+let filter = __ENV.FILTER || "TRUE";
+
+export let options = {
+  stages: [
+    { duration: "5s", target: target },
+    { duration: "60s", target: target },
+  ],
+};
+
+function getQuery(filter, action) {
+  return {
+    action: getAction(action),
+    filterExpression: getFilter(filter),
+  };
+}
+
+function getAction(action) {
+  switch (action) {
+    case "AGGREGATED":
+      return {
+        type: "Aggregated",
+      };
+    case "MUTATIONS":
+      return {
+        minProportion: 0.001,
+        type: "Mutations",
+      };
+    default:
+      return {
+        type: "Aggregated",
+      };
+  }
+}
+
+function getFilter(filter) {
+  switch (filter) {
+    case "TRUE":
+      return {
+        type: "True",
+      };
+    case "INT_EQUALS":
+      return {
+        type: "IntEquals",
+        column: "age",
+        value: 25,
+      };
+    case "COMPLEX_MUTATION":
+      return {
+        type: "And",
+        children: [
+          {
+            children: [
+              {
+                numberOfMatchers: 2,
+                matchExactly: false,
+                children: [
+                  {
+                    sequenceName: "S",
+                    position: 621,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 622,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 624,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 625,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 626,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 627,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 628,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 631,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 634,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 638,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 639,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 640,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 641,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 642,
+                    symbol: "G",
+                    type: "AminoAcidEquals",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 643,
+                    type: "HasAminoAcidMutation",
+                  },
+                  {
+                    sequenceName: "S",
+                    position: 643,
+                    symbol: "-",
+                    type: "AminoAcidEquals",
+                  },
+                ],
+                type: "N-Of",
+              },
+              {
+                child: { position: 897, symbol: "A", type: "NucleotideEquals" },
+                type: "Not",
+              },
+              {
+                child: {
+                  position: 28256,
+                  symbol: "T",
+                  type: "NucleotideEquals",
+                },
+                type: "Not",
+              },
+              {
+                child: {
+                  sequenceName: "ORF1a",
+                  position: 10,
+                  symbol: "*",
+                  type: "AminoAcidEquals",
+                },
+                type: "Not",
+              },
+            ],
+            type: "And",
+          },
+          {
+            type: "DateBetween",
+            column: "date",
+            from: "2021-03-18",
+            to: "2021-03-18",
+          },
+        ],
+      };
+    default:
+      return {
+        type: "True",
+      };
+  }
+}
+
+export default function () {
+  const url = `${host}/query`;
+
+  const query = getQuery(filter, action);
+
+  const params = {
+    headers: {
+      "Content-Type": "application/json",
+    },
+  };
+
+  let response = http.post(url, JSON.stringify(query), params);
+
+  check(response, {
+    "is status 200": (r) => r.status === 200,
+    "log response on error": (r) => {
+      if (r.status !== 200) {
+        console.error(
+          `Request to ${url} failed. Status: ${r.status}, Body: ${r.body}`,
+        );
+      }
+      return r.status === 200;
+    },
+  });
+
+  sleep(1);
+}

loadTests/loadtest.py

@@ -0,0 +1,101 @@
+import argparse
+import logging
+import os
+import subprocess
+from dataclasses import dataclass
+from enum import Enum
+
+logging.basicConfig(level=logging.INFO)
+
+
+class TestActions(Enum):
+    AGGREGATED = 1
+    MUTATIONS = 2
+
+
+class TestFilters(Enum):
+    TRUE = 1
+    INT_EQUALS = 2
+    COMPLEX_MUTATION = 3
+
+
+@dataclass
+class ParsedArgs:
+    action: TestActions
+    filter: TestFilters
+    url: str
+    workers: int
+
+
+def parse_args() -> ParsedArgs:
+    parser = argparse.ArgumentParser(description="SILO load test")
+
+    parser.add_argument(
+        '--url',
+        type=str,
+        default='http://localhost:8091',
+        help='URL to SILO'
+    )
+
+    parser.add_argument('--workers', type=int, default=1, help='Number of workers sending requests')
+    parser.add_argument('--action', type=str, default='AGGREGATED', help='Query to send to SILO')
+    parser.add_argument('--filter', type=str, default='TRUE', help='Filter for query')
+
+    args = parser.parse_args()
+
+    def action_type(arg: str) -> TestActions:
+        try:
+            return TestActions[arg.upper()]
+        except KeyError:
+            raise ValueError(f"Invalid action: {arg}. Valid actions are: AGGREGATED, MUTATIONS")
+
+    def filter_type(arg: str) -> TestFilters:
+        try:
+            return TestFilters[arg.upper()]
+        except KeyError:
+            raise ValueError(f"Invalid filter: {arg}. Valid filters are: TRUE, INT_EQUALS, COMPLEX_MUTATION")
+
+    return ParsedArgs(
+        action=action_type(args.action),
+        filter=filter_type(args.filter),
+        url=args.url,
+        workers=args.workers
+    )
+
+
+def run_load_test(args: ParsedArgs):
+    logging.info(f"Running load test with args: {args}")
+
+    container_name = f"k6_load_test"
+
+    working_directory_of_script = os.path.dirname(os.path.realpath(__file__))
+
+    environment_variables = f"-e TARGET={args.workers} -e HOST={args.url} -e ACTION={args.action.name} -e FILTER={args.filter.name}"
+
+    docker_command = f"docker run --network='host' --name={container_name} -v {working_directory_of_script}:/scripts loadimpact/k6 run /scripts/loadtest.js {environment_variables}"
+
+    logging.info(f"Executing k6 load test in Docker with command: {docker_command}")
+
+    process = subprocess.Popen(docker_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                               preexec_fn=os.setsid)
+
+    try:
+        stdout, stderr = process.communicate()
+        if process.returncode == 0:
+            logging.info("k6 load test in Docker executed successfully")
+            logging.info(stdout.decode())
+        else:
+            logging.error("Error executing k6 load test in Docker")
+            logging.error(stderr.decode())
+    except KeyboardInterrupt:
+        stop_command = f"docker stop {container_name}"
+        subprocess.call(stop_command.split())
+        logging.info("Load test interrupted by user and container stopped")
+
+
+if __name__ == '__main__':
+    try:
+        args = parse_args()
+        run_load_test(args)
+    except KeyboardInterrupt:
+        logging.info("Script execution interrupted by user")