better mongodb handling

1 files changed, 38 insertions, 0 deletions

diff --git a/src/collector/healthcheck.py b/src/collector/healthcheck.py
new file mode 100644
index 0000000..7d336fd
--- /dev/null
+++ b/src/collector/healthcheck.py
@@ -0,0 +1,38 @@
+"""
+Send a healthcheck request
+"""
+import sys
+import time
+import json
+
+import requests
+
+
+def check_collector() -> bool:
+    """Check our collector using /info
+
+    :return: bool
+    """
+    time.sleep(2)  # Prevent race condition with redis container healthcheck
+
+    req = requests.get(
+        "http://localhost:8000/info",
+        timeout=3,
+        # verify="./rootCA.crt",
+    )
+
+    if req.status_code != 200:
+        return False
+
+    data = json.loads(req.text)
+    if isinstance(data["estimated document count"], int) and data["estimated document count"] >= 0:
+        return req.status_code == 200
+
+    return False
+
+
+if __name__ == "__main__":
+    if sys.argv[1] == "COLLECTOR":
+        if check_collector():
+            sys.exit(0)
+    sys.exit(1)