fix(bool_env): fixing and cast env variables

* Fixing case when environment variables are set * Fixing flake8 issues
pryorda · Mar 23, 2019 · 3e9fd0e · 3e9fd0e
1 parent 8adf069
commit 3e9fd0e
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 14 deletions.
diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py
@@ -1,8 +1,10 @@
+import os
+
 from unittest import mock
 
 from pyVmomi import vim
 
-from vmware_exporter.helpers import batch_fetch_properties
+from vmware_exporter.helpers import batch_fetch_properties, get_bool_env
 
 
 class FakeView(vim.ManagedObject):
@@ -14,6 +16,22 @@ def Destroy(self):
         pass
 
 
+def test_get_bool_env_with_default_value():
+    value = get_bool_env('INEXISTENT_ENV', True)
+
+    assert value
+
+
+def test_get_bool_env_with_a_valid_env():
+    key = "TEST_BOOLEAN_VALUE"
+
+    os.environ[key] = "True"
+
+    value = get_bool_env(key, False)
+
+    assert value
+
+
 def test_batch_fetch_properties():
     content = mock.Mock()
 

diff --git a/vmware_exporter/helpers.py b/vmware_exporter/helpers.py
@@ -1,6 +1,12 @@
+import os
+
 from pyVmomi import vmodl
 
 
+def get_bool_env(key, default=None):
+    return bool(os.environ.get(key, default))
+
+
 def batch_fetch_properties(content, obj_type, properties):
     view_ref = content.viewManager.CreateContainerView(
         container=content.rootFolder,

diff --git a/vmware_exporter/vmware_exporter.py b/vmware_exporter/vmware_exporter.py
@@ -31,7 +31,7 @@
 from prometheus_client.core import GaugeMetricFamily
 from prometheus_client import CollectorRegistry, generate_latest
 
-from .helpers import batch_fetch_properties
+from .helpers import batch_fetch_properties, get_bool_env
 from .defer import parallelize, run_once_property
 
 
@@ -762,13 +762,13 @@ def configure(self, args):
                 'vsphere_host': os.environ.get('VSPHERE_HOST'),
                 'vsphere_user': os.environ.get('VSPHERE_USER'),
                 'vsphere_password': os.environ.get('VSPHERE_PASSWORD'),
-                'ignore_ssl': os.environ.get('VSPHERE_IGNORE_SSL', False),
+                'ignore_ssl': get_bool_env('VSPHERE_IGNORE_SSL', False),
                 'collect_only': {
-                    'vms': os.environ.get('VSPHERE_COLLECT_VMS', True),
-                    'vmguests': os.environ.get('VSPHERE_COLLECT_VMGUESTS', True),
-                    'datastores': os.environ.get('VSPHERE_COLLECT_DATASTORES', True),
-                    'hosts': os.environ.get('VSPHERE_COLLECT_HOSTS', True),
-                    'snapshots': os.environ.get('VSPHERE_COLLECT_SNAPSHOTS', True),
+                    'vms': get_bool_env('VSPHERE_COLLECT_VMS', True),
+                    'vmguests': get_bool_env('VSPHERE_COLLECT_VMGUESTS', True),
+                    'datastores': get_bool_env('VSPHERE_COLLECT_DATASTORES', True),
+                    'hosts': get_bool_env('VSPHERE_COLLECT_HOSTS', True),
+                    'snapshots': get_bool_env('VSPHERE_COLLECT_SNAPSHOTS', True),
                 }
             }
         }
@@ -785,13 +785,13 @@ def configure(self, args):
                 'vsphere_host': os.environ.get('VSPHERE_{}_HOST'.format(section)),
                 'vsphere_user': os.environ.get('VSPHERE_{}_USER'.format(section)),
                 'vsphere_password': os.environ.get('VSPHERE_{}_PASSWORD'.format(section)),
-                'ignore_ssl': os.environ.get('VSPHERE_{}_IGNORE_SSL'.format(section), False),
+                'ignore_ssl': get_bool_env('VSPHERE_{}_IGNORE_SSL'.format(section), False),
                 'collect_only': {
-                    'vms': os.environ.get('VSPHERE_{}_COLLECT_VMS'.format(section), True),
-                    'vmguests': os.environ.get('VSPHERE_{}_COLLECT_VMGUESTS'.format(section), True),
-                    'datastores': os.environ.get('VSPHERE_{}_COLLECT_DATASTORES'.format(section), True),
-                    'hosts': os.environ.get('VSPHERE_{}_COLLECT_HOSTS'.format(section), True),
-                    'snapshots': os.environ.get('VSPHERE_{}_COLLECT_SNAPSHOTS'.format(section), True),
+                    'vms': get_bool_env('VSPHERE_{}_COLLECT_VMS'.format(section), True),
+                    'vmguests': get_bool_env('VSPHERE_{}_COLLECT_VMGUESTS'.format(section), True),
+                    'datastores': get_bool_env('VSPHERE_{}_COLLECT_DATASTORES'.format(section), True),
+                    'hosts': get_bool_env('VSPHERE_{}_COLLECT_HOSTS'.format(section), True),
+                    'snapshots': get_bool_env('VSPHERE_{}_COLLECT_SNAPSHOTS'.format(section), True),
                 }
             }