From 438adb44977e27e7444b0ca954345c4c196b90fd Mon Sep 17 00:00:00 2001
From: Tyler Finethy <tyler.finethy@datadoghq.com>
Date: Tue, 28 Jan 2025 14:21:47 -0500
Subject: [PATCH] add: tests for probe budgets

---
 .../probes/probe_snapshot_log_line_budgets.json  | 16 ++++++++++++++++
 tests/debugger/test_debugger_probe_snapshot.py   | 11 +++++++++++
 .../docker/python/flask/debugger_controller.py   |  6 ++++++
 3 files changed, 33 insertions(+)
 create mode 100644 tests/debugger/probes/probe_snapshot_log_line_budgets.json

diff --git a/tests/debugger/probes/probe_snapshot_log_line_budgets.json b/tests/debugger/probes/probe_snapshot_log_line_budgets.json
new file mode 100644
index 0000000000..ff648faecb
--- /dev/null
+++ b/tests/debugger/probes/probe_snapshot_log_line_budgets.json
@@ -0,0 +1,16 @@
+[
+    {
+        "language": "",
+        "type": "",
+        "id": "log170aa-acda-4453-9111-1478a697line",
+        "version": 0,
+        "where": {
+            "typeName": null,
+            "sourceFile": "ACTUAL_SOURCE_FILE",
+            "lines": [
+                "141"
+            ]
+        }
+    }
+    
+]
diff --git a/tests/debugger/test_debugger_probe_snapshot.py b/tests/debugger/test_debugger_probe_snapshot.py
index 83d53487e8..3820ab3d18 100644
--- a/tests/debugger/test_debugger_probe_snapshot.py
+++ b/tests/debugger/test_debugger_probe_snapshot.py
@@ -132,3 +132,14 @@ def test_code_origin_entry_present(self):
                 code_origins_entry_found = code_origin_type == "entry"
 
         assert code_origins_entry_found
+
+    def setup_log_line_probe_snaphots_budgets(self):
+        self._setup("probe_snapshot_log_line_budgets", "/debugger/budgets/5000")
+
+    @missing_feature(context.library != "python", reason="Not yet implemented")
+    def test_log_line_probe_snaphots_budgets(self):
+        self._assert()
+        self._validate_snapshots()
+
+        for _id in self.probe_ids:
+            assert len(self.probe_snapshots[_id]) <= 100
diff --git a/utils/build/docker/python/flask/debugger_controller.py b/utils/build/docker/python/flask/debugger_controller.py
index b6160daf7f..cc568db397 100644
--- a/utils/build/docker/python/flask/debugger_controller.py
+++ b/utils/build/docker/python/flask/debugger_controller.py
@@ -134,3 +134,9 @@ def nulls():
         pii = Pii()
 
     return f"Pii is null {pii is None}. intValue is null {intValue is None}. strValue is null {strValue is None}."
+
+@debugger_blueprint.route("/budgets/<int:loops>", methods=["GET"])
+def budgets(loops):
+    for _ in range(loops):
+        pass
+    return "Budgets", 200