From 6b727bbd6e8f7aec71e84b08d4dd877adf9e6cdf Mon Sep 17 00:00:00 2001
From: Tim Schwenke <tim@trallnag.com>
Date: Wed, 8 Mar 2023 18:10:14 +0100
Subject: [PATCH] feat: Make add() accept args funcs (#230)

* feat: Make add() accept args funcs

* chore: Update CHANGELOG.md
---
 CHANGELOG.md                                  |  7 +++
 .../instrumentation.py                        | 25 ++++----
 tests/test_instrumentator_multiple_apps.py    | 63 ++++++++++++++++++-
 3 files changed, 81 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc2c8fe..9c3bfc9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,13 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0).
 
 ## Unreleased
 
+### Added
+
+- Adjusted the `add()` method to accept an arbitrary number of instrumentation
+  functions as arguments instead of a single one. Non-breaking change.
+  Implemented in pull request
+  [#230](https://github.com/trallnag/prometheus-fastapi-instrumentator/pull/230).
+
 ### Fixed
 
 - Fixed multi process mode in `expose()` method that handles the `/metrics`
diff --git a/src/prometheus_fastapi_instrumentator/instrumentation.py b/src/prometheus_fastapi_instrumentator/instrumentation.py
index 82a75d5..1b2deb5 100644
--- a/src/prometheus_fastapi_instrumentator/instrumentation.py
+++ b/src/prometheus_fastapi_instrumentator/instrumentation.py
@@ -271,7 +271,7 @@ def metrics(request: Request):
 
     def add(
         self,
-        instrumentation_function: Optional[
+        *instrumentation_function: Optional[
             Callable[[metrics.Info], Union[None, Awaitable[None]]]
         ],
     ):
@@ -287,18 +287,19 @@ def add(
             self: Instrumentator. Builder Pattern.
         """
 
-        if instrumentation_function:
-            if asyncio.iscoroutinefunction(instrumentation_function):
-                self.async_instrumentations.append(
-                    cast(
-                        Callable[[metrics.Info], Awaitable[None]],
-                        instrumentation_function,
+        for func in instrumentation_function:
+            if func:
+                if asyncio.iscoroutinefunction(func):
+                    self.async_instrumentations.append(
+                        cast(
+                            Callable[[metrics.Info], Awaitable[None]],
+                            func,
+                        )
+                    )
+                else:
+                    self.instrumentations.append(
+                        cast(Callable[[metrics.Info], None], func)
                     )
-                )
-            else:
-                self.instrumentations.append(
-                    cast(Callable[[metrics.Info], None], instrumentation_function)
-                )
 
         return self
 
diff --git a/tests/test_instrumentator_multiple_apps.py b/tests/test_instrumentator_multiple_apps.py
index 6bbae14..98241d7 100644
--- a/tests/test_instrumentator_multiple_apps.py
+++ b/tests/test_instrumentator_multiple_apps.py
@@ -2,7 +2,7 @@
 from prometheus_client import CollectorRegistry, Counter
 from starlette.testclient import TestClient
 
-from prometheus_fastapi_instrumentator import Instrumentator
+from prometheus_fastapi_instrumentator import Instrumentator, metrics
 
 
 def test_multiple_apps_custom_registry():
@@ -63,7 +63,7 @@ async def dummy_middleware(request, call_next):
 
 
 def test_multiple_apps_expose_defaults():
-    """Tests instrumentation of multiple apps in combination with  middlewares."""
+    """Tests instrumentation of multiple apps in combination with middlewares."""
 
     app1 = FastAPI()
     app2 = FastAPI()
@@ -104,3 +104,62 @@ async def dummy_middleware(request, call_next):
     want = 'http_requests_total{handler="/dummy",method="GET",status="2xx"} 1.0\n'
     assert want in metrics1
     assert want in metrics2
+
+
+def test_multiple_apps_expose_full():
+    """Tests instrumentation of multiple apps in combination with middlewares."""
+
+    app1 = FastAPI()
+    app2 = FastAPI()
+
+    @app1.get("/dummy")
+    def read_dummy_app1():
+        return "Hello from app1"
+
+    @app2.get("/dummy")
+    def read_dummy_app2():
+        return "Hello from app2"
+
+    Instrumentator().add(
+        metrics.request_size(),
+        metrics.requests(),
+        metrics.combined_size(),
+        metrics.response_size(),
+        metrics.latency(),
+        metrics.default(),
+    ).instrument(app1).expose(app1)
+
+    Instrumentator().add(
+        metrics.request_size(),
+        metrics.requests(),
+        metrics.combined_size(),
+        metrics.response_size(),
+        metrics.latency(),
+        metrics.default(),
+    ).instrument(app2).expose(app2)
+
+    # Add middleware after adding the instrumentator, this triggers another
+    # app.build_middleware_stack(), which creates the middleware again, but it
+    # will use the same Prometheus registry again, which could try to create the
+    # same metrics again causing duplication errors.
+    @app1.middleware("http")
+    @app2.middleware("http")
+    async def dummy_middleware(request, call_next):
+        response = await call_next(request)
+        return response
+
+    client1 = TestClient(app1)
+    client2 = TestClient(app2)
+
+    client1.get("/dummy")
+    client2.get("/dummy")
+
+    metrics1 = client1.get("/metrics").content.decode()
+    metrics2 = client2.get("/metrics").content.decode()
+
+    print("app1 GET /metrics\n" + metrics1)
+    print("app2 GET /metrics\n" + metrics2)
+
+    want = 'http_requests_total{handler="/dummy",method="GET",status="2xx"} 1.0\n'
+    assert want in metrics1
+    assert want in metrics2