ci(perf): Add workflow concurrency stress tests

.github/workflows/benchmark.yml

@@ -0,0 +1,67 @@
+name: Benchmark
+
+on:
+  push:
+    branches: ["main"]
+    paths:
+      - "docker-compose.yml"
+      - "Dockerfile"
+      - "pyproject.toml"
+      - "tests/**"
+      - "tracecat-cli/**"
+      - "tracecat/**"
+  pull_request:
+    branches: ["main"]
+    paths: [".github/workflows/benchmark.yml"]
+
+jobs:
+  stress-test:
+    runs-on: ubuntu-latest-8-cores
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        num_workers: ["1", "2", "4", "8"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.12
+        uses: actions/setup-python@v3
+        with:
+          python-version: "3.12"
+          cache: "pip"
+          cache-dependency-path: pyproject.toml
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Run environment setup script
+        run: bash env.sh
+
+      - name: Get image tag
+        id: set-image-tag
+        run: |
+          if [ "${{ github.event_name }}" == "push" ] && [ "${{ github.ref }}" == "refs/heads/main" ]; then
+            echo "TRACECAT__IMAGE_TAG=main" >> $GITHUB_ENV
+          else
+            echo "TRACECAT__IMAGE_TAG=pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
+          fi
+
+      - name: Start Docker services
+        env:
+          NUM_WORKERS: ${{ matrix.num_workers }}
+        run: |
+          docker compose up --no-deps api worker postgres_db temporal -d
+
+      - name: pip install Tracecat
+        run: python -m pip install --upgrade pip && pip install ".[dev]" && pip install ./cli
+
+      - name: Run tests (headless mode)
+        env:
+          TRACECAT__IMAGE_TAG: ${{ env.TRACECAT__IMAGE_TAG }}
+          LOG_LEVEL: WARNING
+        run: |
+          pytest -k "test_stress" --temporal-no-restart --tracecat-no-restart
+              --benchmark-name=short \
+              --benchmark-group-by=param \
+              --benchmark-warmup=off \
+              --benchmark-columns=min,max,mean,median,stddev,iterations \

.github/workflows/build-push-images.yml

@@ -2,8 +2,6 @@ name: Publish Images
 
 on:
   push:
-    branches:
-      - "main"
     tags:
       - "*.*.*"
 

.github/workflows/test.yml

@@ -3,16 +3,22 @@ name: Tests
 on:
   push:
     branches: ["main"]
-    paths-ignore:
-      - "frontend/**"
-      - "docs/**"
-      - "README.md"
+    paths:
+      - "docker-compose.yml"
+      - "Dockerfile"
+      - "pyproject.toml"
+      - "tests/**"
+      - "tracecat-cli/**"
+      - "tracecat/**"
   pull_request:
     branches: ["main"]
-    paths-ignore:
-      - "frontend/**"
-      - "docs/**"
-      - "README.md"
+    paths:
+      - "docker-compose.yml"
+      - "Dockerfile"
+      - "pyproject.toml"
+      - "tests/**"
+      - "tracecat-cli/**"
+      - "tracecat/**"
 
 permissions:
   contents: read
@@ -100,9 +106,6 @@ jobs:
       - name: Start Docker services
         run: docker compose up --no-deps api worker postgres_db -d
 
-      - name: Verify Tracecat API is running
-        run: curl -s http://localhost:8000/health | jq -e '.status == "ok"'
-
       - name: pip install Tracecat
         run: python -m pip install --upgrade pip && pip install ".[dev]" && pip install ./cli
 
@@ -154,9 +157,6 @@ jobs:
       - name: Start Docker services
         run: docker compose up --no-deps api worker postgres_db temporal -d
 
-      - name: Verify Tracecat API is running
-        run: curl -s http://localhost:8000/health | jq -e '.status == "ok"'
-
       - name: pip install Tracecat
         run: python -m pip install --upgrade pip && pip install ".[dev]" && pip install ./cli
 
@@ -202,4 +202,4 @@ jobs:
           URLSCAN_API_KEY: ${{ secrets.INTEGRATION__URLSCAN_API_KEY }}
           VT_API_KEY: ${{ secrets.INTEGRATION__VT_API_KEY }}
           LOG_LEVEL: WARNING
-        run: pytest -k "test_playbooks" --temporal-no-restart --tracecat-no-restart
+        run: pytest -k "test_playbooks" --benchmark-disable --temporal-no-restart --tracecat-no-restart

pyproject.toml

@@ -67,6 +67,7 @@ Repository = "https://github.com/TracecatHQ/tracecat"
 dev = [
     "respx",
     "pytest",
+    "pytest-benchmark",
     "python-dotenv",
     "pytest-asyncio",
     "pytest-mock==3.14.0",

tests/conftest.py

@@ -1,12 +1,19 @@
+import asyncio
 import os
 import subprocess
 import time
+from pathlib import Path
 from uuid import uuid4
 
 import pytest
+import yaml
 from cryptography.fernet import Fernet
 from loguru import logger
 
+from tracecat.dsl.common import DSLInput
+
+DATA_PATH = Path(__file__).parent.joinpath("data/workflows")
+
 
 def pytest_addoption(parser: pytest.Parser):
     parser.addoption(
@@ -171,3 +178,73 @@ def tracecat_worker(env_sandbox):
             ["docker", "compose", "down", "--remove-orphans", "worker"], check=True
         )
         logger.info("Stopped Tracecat Temporal worker")
+
+
+@pytest.fixture
+def mock_registry():
+    """Mock registry for testing UDFs.
+
+    Note
+    ----
+    - This fixture is used to test the integration of UDFs with the workflow.
+    - It's unreachable by an external worker, as the worker will not have access
+    to these functions when it starts up.
+    """
+    from tracecat.registry import registry
+
+    # NOTE!!!!!!!: Didn't want to spend too much time figuring out how
+    # to grab the actual execution order using the client, so I'm using a
+    # hacky way to get the order of execution. TO FIX LATER
+    # The counter doesn't get reset properly so you should never use this outside
+    # of the 'ordering' tests
+    def counter():
+        i = 0
+        while True:
+            yield i
+            i += 1
+
+    counter_gen = counter()
+    if "integration_test.count" not in registry:
+
+        @registry.register(
+            description="Counts up from 0",
+            namespace="integration_test",
+        )
+        def count(arg: str | None = None) -> int:
+            order = next(counter_gen)
+            return order
+
+    if "integration_test.passthrough" not in registry:
+
+        @registry.register(
+            description="passes through",
+            namespace="integration_test",
+        )
+        async def passthrough(num: int) -> int:
+            await asyncio.sleep(0.1)
+            return num
+
+    registry.init()
+    yield registry
+    counter_gen = counter()  # Reset the counter generator
+
+
+@pytest.fixture
+def dsl(request: pytest.FixtureRequest) -> DSLInput:
+    test_name = request.param
+    data_path = DATA_PATH / f"{test_name}.yml"
+    dsl = DSLInput.from_yaml(data_path)
+    return dsl
+
+
+@pytest.fixture
+def dsl_with_expected(request: pytest.FixtureRequest) -> DSLInput:
+    test_name = request.param
+    data_path = DATA_PATH / f"{test_name}.yml"
+    expected_path = DATA_PATH / f"{test_name}_expected.yml"
+    dsl = DSLInput.from_yaml(data_path)
+    with expected_path.open() as f:
+        yaml_data = f.read()
+    data = yaml.safe_load(yaml_data)
+    expected_result = {key: (value or {}) for key, value in data.items()}
+    return dsl, expected_result

tests/integration/test_stress.py

@@ -0,0 +1,37 @@
+import asyncio
+import uuid
+from pathlib import Path
+
+import pytest
+
+from tracecat.dsl.dispatcher import dispatch_workflow
+
+DATA_PATH = Path(__file__).parent.parent.joinpath("data/workflows")
+TEST_WF_ID = "wf-00000000000000000000000000000000"
+
+
+@pytest.mark.parametrize(
+    "dsl",
+    ["stress_adder_tree"],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "num_workflows", [10, 100, 1000], ids=lambda x: f"num_workflows={x}"
+)
+@pytest.mark.slow
+def test_concurrent_workflows(
+    dsl, num_workflows, temporal_cluster, mock_registry, auth_sandbox, benchmark
+):
+    """Multiple executions of the same workflow run at the same time."""
+
+    def generate_wf_id():
+        return f"wf-{uuid.uuid4()}"
+
+    async def run_workflows():
+        tasks = [
+            dispatch_workflow(dsl=dsl, wf_id=generate_wf_id())
+            for _ in range(num_workflows)
+        ]
+        return await asyncio.gather(*tasks)
+
+    benchmark.pedantic(lambda: asyncio.run(run_workflows()), iterations=3, rounds=1)

tests/integration/test_validation.py

@@ -20,7 +20,6 @@ def filename(request: pytest.FixtureRequest) -> Path:
 )
 @pytest.mark.asyncio
 async def test_workflow_commit(filename, auth_sandbox):
-    print(filename)
     title = f"Test workflow: {filename}"
     workflow_result = await shared.create_workflow(title)
     await shared.commit_workflow(filename, workflow_result["id"])