add integration test for chain relays

getsentry · Dav1dde · Mar 19, 2024 · Mar 18, 2024 · Mar 19, 2024 · Mar 19, 2024
commit 485ead9f457405675a56d61bd8897982e81b88e4
@@ -13,7 +13,12 @@
 from .fixtures.haproxy import haproxy  # noqa
 from .fixtures.mini_sentry import mini_sentry  # noqa
 from .fixtures.aws_lambda_runtime import aws_lambda_runtime  # noqa
-from .fixtures.relay import relay, get_relay_binary, latest_relay_version  # noqa
+from .fixtures.relay import (  # noqa
+    relay,
+    relay_credentials,
+    get_relay_binary,
+    latest_relay_version,
+)
 from .fixtures.processing import (  # noqa
     kafka_consumer,
     get_topic_name,

@@ -80,9 +80,9 @@ def relay_with_processing(relay, mini_sentry, processing_config):
     requests to the test ingestion topics
     """
 
-    def inner(options=None):
+    def inner(options=None, **kwargs):
         options = processing_config(options)
-        return relay(mini_sentry, options=options)
+        return relay(mini_sentry, options=options, **kwargs)
 
     return inner
 

@@ -1,6 +1,7 @@
 import json
 import os
 import sys
+import uuid
 import signal
 import stat
 import requests
@@ -9,6 +10,7 @@
 import yaml
 import pytest
 
+from sentry_relay.auth import generate_key_pair
 from . import SentryLike
 
 RELAY_BIN = [os.path.abspath(os.environ.get("RELAY_BIN") or "target/debug/relay")]
@@ -90,6 +92,15 @@ def inner(version="latest"):
     return inner
 
 
+@pytest.fixture
+def relay_credentials():
+    def inner():
+        sk, pk = generate_key_pair()
+        return {"public_key": str(pk), "secret_key": str(sk), "id": str(uuid.uuid4())}
+
+    return inner
+
+
 @pytest.fixture
 def relay(mini_sentry, random_port, background_process, config_dir, get_relay_binary):
     def inner(
@@ -99,6 +110,7 @@ def inner(
         external=None,
         wait_health_check=True,
         static_relays=None,
+        credentials=None,
         version="latest",
     ):
         relay_bin = get_relay_binary(version)
@@ -146,11 +158,16 @@ def inner(
         dir = config_dir("relay")
         dir.join("config.yml").write(yaml.dump(default_opts))
 
-        subprocess.check_output(relay_bin + ["-c", str(dir), "credentials", "generate"])
+        if credentials is None:
+            subprocess.check_output(
+                relay_bin + ["-c", str(dir), "credentials", "generate"]
+            )
+            with open(dir.join("credentials.json")) as f:
+                credentials = json.load(f)
+        else:
+            with open(dir.join("credentials.json"), "w") as f:
+                f.write(json.dumps(credentials))
 
-        # now that we have generated a credentials file get the details
-        with open(dir.join("credentials.json")) as f:
-            credentials = json.load(f)
         public_key = credentials.get("public_key")
         assert public_key is not None
         secret_key = credentials.get("secret_key")

@@ -1,4 +1,5 @@
 from typing import Any
+import pytest
 from attr import dataclass
 from .test_metrics import TEST_CONFIG
 
@@ -24,11 +25,29 @@ def metric_stats_by_mri(metrics_consumer, count, timeout=None):
     return MetricStatsByMri(volume=volume, other=other)
 
 
-def test_metric_stats_simple(mini_sentry, relay_with_processing, metrics_consumer):
+@pytest.mark.parametrize("mode", ["default", "chain"])
+def test_metric_stats_simple(
+    mini_sentry, relay, relay_with_processing, relay_credentials, metrics_consumer, mode
+):
     mini_sentry.global_config["options"]["relay.metric-stats.rollout-rate"] = 1.0
 
     metrics_consumer = metrics_consumer()
-    relay = relay_with_processing(options=TEST_CONFIG)
+
+    if mode == "default":
+        relay = relay_with_processing(options=TEST_CONFIG)
+    elif mode == "chain":
+        credentials = relay_credentials()
+        static_relays = {
+            credentials["id"]: {
+                "public_key": credentials["public_key"],
+                "internal": True,
+            },
+        }
+        relay = relay(
+            relay_with_processing(options=TEST_CONFIG, static_relays=static_relays),
+            options=TEST_CONFIG,
+            credentials=credentials,
+        )
 
     project_id = 42
     project_config = mini_sentry.add_basic_project_config(project_id)