From b5d362417d7415fafe9900b7f10f5c59f16dd900 Mon Sep 17 00:00:00 2001
From: Mathias Petermann <mathias.petermann@gmail.com>
Date: Thu, 3 Mar 2022 09:08:27 +0100
Subject: [PATCH] feat: add filter option exclude_tags (#159)

---
 docs/content/configuration/defaults.md      |  3 +++
 prometheuspvesd/config.py                   |  6 ++++++
 prometheuspvesd/discovery.py                |  6 ++++++
 prometheuspvesd/test/fixtures/fixtures.py   | 22 ++++++++++++++++++++-
 prometheuspvesd/test/unit/test_discovery.py | 22 +++++++++++++++++++--
 5 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/docs/content/configuration/defaults.md b/docs/content/configuration/defaults.md
index 27e27ef..1419201 100644
--- a/docs/content/configuration/defaults.md
+++ b/docs/content/configuration/defaults.md
@@ -28,6 +28,9 @@ exclude_state: []
 # needs to be a list of strings
 exclude_vmid: []
 
+# can be used to exclude vms by tags (proxmox 6+)
+exclude_tags: []
+
 pve:
     server:
     user:
diff --git a/prometheuspvesd/config.py b/prometheuspvesd/config.py
index de975d8..4b17387 100644
--- a/prometheuspvesd/config.py
+++ b/prometheuspvesd/config.py
@@ -94,6 +94,12 @@ class Config():
             "file": True,
             "type": environs.Env().list
         },
+        "exclude_tags": {
+            "default": [],
+            "env": "EXCLUDE_TAGS",
+            "file": True,
+            "type": environs.Env().list
+        },
         "pve.server": {
             "default": "",
             "env": "PVE_SERVER",
diff --git a/prometheuspvesd/discovery.py b/prometheuspvesd/discovery.py
index 32f607f..5a2b471 100644
--- a/prometheuspvesd/discovery.py
+++ b/prometheuspvesd/discovery.py
@@ -157,6 +157,12 @@ class Discovery():
             if str(obj["vmid"]) in self.config.config["exclude_vmid"]:
                 continue
 
+            if (
+                isinstance(obj["tags"], str)
+                and not set(obj["tags"].split(",")).isdisjoint(self.config.config["exclude_tags"])
+            ):
+                continue
+
             filtered.append(item.copy())
         return filtered
 
diff --git a/prometheuspvesd/test/fixtures/fixtures.py b/prometheuspvesd/test/fixtures/fixtures.py
index ca17c29..23b0c96 100644
--- a/prometheuspvesd/test/fixtures/fixtures.py
+++ b/prometheuspvesd/test/fixtures/fixtures.py
@@ -34,7 +34,8 @@ def qemus():
             "disk": 0,
             "status": "running",
             "netout": 12159205236,
-            "mem": 496179157
+            "mem": 496179157,
+            "tags": "unmonitored,excluded"
         },
         {
             "diskwrite": 0,
@@ -54,6 +55,25 @@ def qemus():
             "netout": 12159205236,
             "mem": 496179157
         },
+        {
+            "diskwrite": 0,
+            "vmid": "102",
+            "name": "102.example.com",
+            "cpu": 0.0202130478509556,
+            "diskread": 0,
+            "template": "",
+            "uptime": 3101505,
+            "maxdisk": 26843545600,
+            "maxmem": 1073741824,
+            "pid": "1765",
+            "cpus": 1,
+            "netin": 2856071643,
+            "disk": 0,
+            "status": "prelaunch",
+            "netout": 12159205236,
+            "mem": 496179157,
+            "tags": "monitored"
+        },
     ]
 
 
diff --git a/prometheuspvesd/test/unit/test_discovery.py b/prometheuspvesd/test/unit/test_discovery.py
index 5397482..a1f4ba4 100644
--- a/prometheuspvesd/test/unit/test_discovery.py
+++ b/prometheuspvesd/test/unit/test_discovery.py
@@ -17,13 +17,31 @@ def get_mock(*args):
     return False
 
 
-def test_exclude(discovery_fixture, qemus):
-    discovery_fixture.config.config["exclude_vmid"] = ["100", "101"]
+def test_exclude_vmid(discovery_fixture, qemus):
+    discovery_fixture.config.config["exclude_vmid"] = ["100", "101", "102"]
 
     expected = []
     filtered = discovery_fixture._exclude(qemus)
 
     assert filtered == expected
+    discovery_fixture.config.config["exclude_vmid"] = []
+
+
+def test_exclude_state(discovery_fixture, qemus):
+    discovery_fixture.config.config["exclude_state"] = ["prelaunch"]
+    filtered = discovery_fixture._exclude(qemus)
+
+    assert len(filtered) == 2
+    discovery_fixture.config.config["exclude_state"] = []
+
+
+def test_exclude_tags(discovery_fixture, qemus):
+    discovery_fixture.config.config["exclude_tags"] = ["unmonitored"]
+
+    filtered = discovery_fixture._exclude(qemus)
+
+    assert len(filtered) == 2
+    discovery_fixture.config.config["exclude_tags"] = []
 
 
 def test_validate_ip(discovery_fixture, addresses):