From 52ccd851fe38f591f50bd7a3e857904a8817d91d Mon Sep 17 00:00:00 2001
From: Mathias Brulatout <m.brulatout@criteo.com>
Date: Fri, 15 Mar 2024 13:22:03 +0100
Subject: [PATCH] feat: Add an extra_checks param to service registration

Without breaking compatibility, an optional extra_checks
list of Checks allows us to register multiple checks during
service registration.

Note that this was still possible through check registration
but requires multiple API calls.
---
 consul/base.py    |  5 ++++-
 tests/test_std.py | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/consul/base.py b/consul/base.py
index 3415a6d..8c581ce 100644
--- a/consul/base.py
+++ b/consul/base.py
@@ -836,6 +836,7 @@ def register(
                 http=None,
                 timeout=None,
                 enable_tag_override=False,
+                extra_checks=None,
             ):
                 """
                 Add a new service to the local agent. There is more
@@ -877,6 +878,8 @@ def register(
                 https://www.consul.io/docs/agent/services.html
                 """
 
+                if extra_checks is None:
+                    extra_checks = []
                 payload = {}
 
                 payload["name"] = name
@@ -893,7 +896,7 @@ def register(
                 if meta:
                     payload["meta"] = meta
                 if check:
-                    payload["check"] = check
+                    payload["checks"] = [check] + extra_checks
                 if weights:
                     payload["weights"] = weights
 
diff --git a/tests/test_std.py b/tests/test_std.py
index b351724..10ae457 100644
--- a/tests/test_std.py
+++ b/tests/test_std.py
@@ -245,6 +245,35 @@ def verify_check_status(check_id, status, notes=None):
         verify_check_status("ttl_check", "critical")
         verify_and_dereg_check("ttl_check")
 
+    def test_service_multi_check(self, consul_port):
+        c = consul.Consul(port=consul_port)
+        http_addr = f"http://127.0.0.1:{consul_port}"
+        c.agent.service.register(
+            "foo1",
+            check=Check.http(http_addr, "10ms"),
+            extra_checks=[
+                Check.http(http_addr, "20ms"),
+                Check.http(http_addr, "30ms"),
+            ],
+        )
+
+        time.sleep(200 / 1000.0)
+
+        _index, nodes = c.health.service("foo1")
+        assert {check["ServiceID"] for node in nodes for check in node["Checks"]} == {"foo1", ""}
+
+        assert {check["CheckID"] for node in nodes for check in node["Checks"]} == {
+            "service:foo1:1",
+            "service:foo1:2",
+            "service:foo1:3",
+            "serfHealth",
+        }
+        time.sleep(1)
+
+        _index, checks = c.health.checks(service="foo1")
+        assert [check["CheckID"] for check in checks] == ["service:foo1:1", "service:foo1:2", "service:foo1:3"]
+        assert [check["Status"] for check in checks] == ["passing", "passing", "passing"]
+
     def test_service_dereg_issue_156(self, consul_port):
         # https://github.com/cablehead/python-consul/issues/156
         service_name = "app#127.0.0.1#3000"
@@ -274,15 +303,17 @@ def test_agent_checks_service_id(self, consul_obj):
         assert [node["Service"]["ID"] for node in nodes] == ["foo1"]
 
         c.agent.check.register("foo", Check.ttl("100ms"), service_id="foo1")
+        c.agent.check.register("foo2", Check.ttl("100ms"), service_id="foo1")
 
         time.sleep(40 / 1000.0)
 
         _index, nodes = c.health.service("foo1")
         assert {check["ServiceID"] for node in nodes for check in node["Checks"]} == {"foo1", ""}
-        assert {check["CheckID"] for node in nodes for check in node["Checks"]} == {"foo", "serfHealth"}
+        assert {check["CheckID"] for node in nodes for check in node["Checks"]} == {"foo", "foo2", "serfHealth"}
 
         # Clean up tasks
         assert c.agent.check.deregister("foo") is True
+        assert c.agent.check.deregister("foo2") is True
 
         time.sleep(40 / 1000.0)