From 3ae7ac7219532558c3d344ec336bc4c0edc45cf1 Mon Sep 17 00:00:00 2001
From: Ricky Moorhouse <ricky@samespirit.net>
Date: Tue, 9 Oct 2018 11:02:14 +0100
Subject: [PATCH] Support for a boolean to include or exlude

In the discovery arguments, set `enabled_key` to a key
containing a boolean value as to whether to include.
If there is no key an object is included
---
 hemApp/drivers/discovery_file.py    |  7 ++++++-
 tests/hosts_key.yaml                |  5 +++++
 tests/test_driver_discovery_file.py | 15 ++++++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/hemApp/drivers/discovery_file.py b/hemApp/drivers/discovery_file.py
index abdeb2b..4b2dcc3 100644
--- a/hemApp/drivers/discovery_file.py
+++ b/hemApp/drivers/discovery_file.py
@@ -13,9 +13,14 @@ def hosts(**kwargs):
     try:
         with open(kwargs['name'], 'rt') as source_file:
             hosts = yaml.safe_load(source_file)
+            print(hosts)
             for host in hosts:
                 if 'key' in kwargs:
-                    results.append(host.get(kwargs['key']))
+                    if 'enabled_key' in kwargs:
+                        if host.get(kwargs['enabled_key'], True) == True:
+                            results.append(host.get(kwargs['key']))
+                    else:
+                        results.append(host.get(kwargs['key']))
                 else:
                     results.append(host)
     except FileNotFoundError:
diff --git a/tests/hosts_key.yaml b/tests/hosts_key.yaml
index 6857da9..126de6a 100644
--- a/tests/hosts_key.yaml
+++ b/tests/hosts_key.yaml
@@ -1,6 +1,11 @@
  - 
    hostname: host1.example.com
    type: web
+   check: true
+ - 
+   hostname: host0.example.com
+   type: web
+   check: false
  - 
    hostname: host2.example.com
    type: web
diff --git a/tests/test_driver_discovery_file.py b/tests/test_driver_discovery_file.py
index a1da2d0..ff5bf32 100644
--- a/tests/test_driver_discovery_file.py
+++ b/tests/test_driver_discovery_file.py
@@ -45,6 +45,19 @@ def test_check_content_key():
     print(hosts)
     assert 'host1.example.com' in hosts
 
+def test_check_content_key_enable():
+    import hemApp
+    hosts = hemApp.discover_hosts({
+        "type":"file",
+        "name":"tests/hosts_key.yaml",
+        "enabled_key":"check",
+        "key":"hostname"})
+    assert type(hosts) == list
+    print(hosts)
+    assert 'host0.example.com' not in hosts
+    assert 'host1.example.com' in hosts
+    assert 'host2.example.com' in hosts
+
 def test_check_metrics(capsys):
     import hemApp
     metrics = hemApp.initialise_metrics({"type":"console"})
@@ -60,4 +73,4 @@ def test_check_metrics(capsys):
         print("Capsys not working in python 3.3 or 3.4 - investigating")
     
 if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+    unittest.main()