"fixed" the unittest :)

oliver-zehentleitner · oliver-zehentleitner · commit c939fd4fdd59 · 2019-06-12T15:41:02.000+02:00
diff --git a/README.md b/README.md
@@ -23,15 +23,15 @@ OK: 27.8 celsius | celsius=27.8
 ```
 
 ## Config files:
-1. copy `check_pcmeasure.py' to your Nagios/Icinga plugin folder
-2. define a command:
+- copy `check_pcmeasure.py' to your Nagios/Icinga plugin folder
+- define a command:
 ```
 define command{
     command_name       check_pcmeasure
     command_line       $USER1$/check_pcmeasure.py -H $HOSTADDRESS$ -w $ARG1$ -c $ARG2$ -S $ARG3 -l $ARG4
 }
 ```
-3. define a service
+- define a service
 ```
 define service{
     use                     generic-service
@@ -41,7 +41,7 @@ define service{
 }
 ```
 ## Director:
-1. Icinga Director -> Define Data Fields:
+- Icinga Director -> Define Data Fields:
 
 |Label|Field name|
 |-----|----------|
@@ -50,15 +50,14 @@ define service{
 |Etherbox address|etherbox_address|
 |Etherbox sensor label|etherbox_sensor_label|
 |Etherbox sensor name|etherbox_sensor_name|
-
-2. Icinga Director -> Commands -> Templates -> Add:
+- Icinga Director -> Commands -> Templates -> Add:
 ![screenshot](https://s3.gifyu.com/images/check_pcmeasure_command_template.png)
-3. Click on "Fields":
+- Click on "Fields":
 ![screenshot](https://s3.gifyu.com/images/check_pcmeasure_command_template_fields.png)
-4. Icinga Director -> Commands -> Commands -> Add:
+- Icinga Director -> Commands -> Commands -> Add:
 ![screenshot](https://s3.gifyu.com/images/check_pcmeasure_command.png)
-5. Click on "Arguments"
+- Click on "Arguments"
 ![screenshot](https://s3.gifyu.com/images/check_pcmeasure_command_arguments.png)
-6. Create a service as usual ...
-
+- Create a service as usual ...
+- Enjoy :)
 ![screenshot](https://s3.gifyu.com/images/perfdata_icinga.png)
diff --git a/check_pcmeasure.py b/check_pcmeasure.py
@@ -17,45 +17,67 @@
 
 
 class CheckPcMeasure(object):
-    def __init__(self):
+    def __init__(self, host="127.0.0.1", port=4000):
         version = "0.2.0"
-        parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
-                                         description=textwrap.dedent(os.path.basename(__file__) + " " + version +
-                                                                     " (GPLv3) 2018 - " + time.strftime("%Y")),
-                                         epilog=textwrap.dedent("GitHub: https://github.com/mpibpc-mroose/"
-                                                                "nagios_plugin_pcmeasure"))
-        parser.add_argument("-H", dest="host", type=str, required=True, help="hostname of the etherbox")
-        parser.add_argument("-p", type=int, dest="port", default=4000, required=False, help="port to use")
-        parser.add_argument("-S", dest="sensor_name", type=str, required=True, help="sensor name, e.g. com1.1")
-        parser.add_argument("-w", dest="warning_threshold", type=float, required=True, help="warning threshold")
-        parser.add_argument("-c", dest="critical_threshold", type=float, required=True, help="critical threshold")
-        parser.add_argument("-l", dest="label", type=str, required=False, help="label for perfdata, e.g. 'celsius'")
-        self.arguments = parser.parse_args()
-
-        if self.arguments.label:
-            self.label_perfdata = self.arguments.label
-            self.arguments.label = " " + self.arguments.label
-        else:
-            self.label_perfdata = "sensor_output"
-            self.arguments.label = ""
+        if __name__ == '__main__':
+            parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
+                                             description=textwrap.dedent(os.path.basename(__file__) + " " + version +
+                                                                         " (GPLv3) 2018 - " + time.strftime("%Y")),
+                                             epilog=textwrap.dedent("GitHub: https://github.com/mpibpc-mroose/"
+                                                                    "nagios_plugin_pcmeasure"))
+            parser.add_argument("-H", dest="host", type=str, required=True, help="hostname of the etherbox")
+            parser.add_argument("-p", type=int, dest="port", default=4000, required=False, help="port to use")
+            parser.add_argument("-S", dest="sensor_name", type=str, required=True, help="sensor name, e.g. com1.1")
+            parser.add_argument("-w", dest="warning_threshold", type=float, required=True, help="warning threshold")
+            parser.add_argument("-c", dest="critical_threshold", type=float, required=True, help="critical threshold")
+            parser.add_argument("-l", dest="label", type=str, required=False, help="label for perfdata, e.g. 'celsius'")
 
-        if self.arguments.warning_threshold > self.arguments.critical_threshold:
-            print("UNKNOWN: warning threshold should be lower than critical threshold")
-            sys.exit(STATUS_UNKNOWN)
+            self.arguments = parser.parse_args()
+            self.critical_threshold = self.arguments.critical_threshold
+            self.warning_threshold = self.arguments.warning_threshold
+            self.host = self.arguments.host
+            self.port = self.arguments.port
+            self.sensor_name = self.arguments.sensor_name
 
+            if self.arguments.label:
+                self.label_name = " " + self.arguments.label
+                self.label_perfdata = self.arguments.label
+            else:
+                self.label_name = ""
+                self.label_perfdata = "sensor_output"
+
+            if self.warning_threshold > self.critical_threshold:
+                print("UNKNOWN: warning threshold should be lower than critical threshold")
+                sys.exit(STATUS_UNKNOWN)
+        else:
+            self.label_name = ""
+            self.label_perfdata = "sensor_output"
+            self.critical_threshold = False
+            self.warning_threshold = False
+            self.host = host
+            self.port = port
+            self.sensor_name = False
         self.socket = self._connect()
 
     def _connect(self):
         _socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         _socket.settimeout(0.5)
         try:
-            _socket.connect((self.arguments.host, self.arguments.port))
+            _socket.connect((self.host, self.port))
         except socket.timeout:
-            print("UNKNOWN: timeout on connect to {host}:{port}".format(host=self.arguments.host,
-                                                                        port=self.arguments.port))
+            print("UNKNOWN: timeout on connect to {host}:{port}".format(host=self.host,
+                                                                        port=self.port))
+            sys.exit(STATUS_UNKNOWN)
+        except ConnectionRefusedError as error:
+            print("UNKNOWN: {error} by {host}:{port}".format(error=str(error),
+                                                             host=self.host,
+                                                             port=self.port))
             sys.exit(STATUS_UNKNOWN)
         return _socket
 
+    def close(self):
+        self.socket.close()
+
     def __delete__(self, instance):
         self.socket.close()
 
@@ -65,38 +87,53 @@ def _get_value(self, sensor_name) -> float:
         sensor_reply_pattern = re.compile(r"value=\s(?P<value>.+);")
         return float(sensor_reply_pattern.findall(raw_data)[0])
 
-    def check(self):
+    def check(self, sensor_name=False, warning_threshold=False, critical_threshold=False, label=False) -> int:
+        if sensor_name:
+            self.sensor_name = sensor_name
+        if warning_threshold:
+            self.warning_threshold = warning_threshold
+        if critical_threshold:
+            self.critical_threshold = critical_threshold
+        if label:
+            self.label_name = " " + label
+            self.label_perfdata = label
         try:
-            value = self._get_value(sensor_name=self.arguments.sensor_name)
-            if value > self.arguments.critical_threshold:
+            value = self._get_value(sensor_name=self.sensor_name)
+            if value > self.critical_threshold:
                 print("CRITICAL: {value} exceeds {threshold} | "
                       "{label_perfdata}={value};;;0".format(value=value,
-                                                            label=self.arguments.label,
+                                                            label=self.label_name,
                                                             label_perfdata=self.label_perfdata,
-                                                            threshold=self.arguments.critical_threshold))
-                sys.exit(STATUS_CRITICAL)
-            elif value > self.arguments.warning_threshold:
+                                                            threshold=self.critical_threshold))
+                return STATUS_CRITICAL
+            elif value > self.warning_threshold:
                 print("WARNING: {value} exceeds {threshold} | "
                       "{label_perfdata}={value};;;0".format(value=value,
-                                                            label=self.arguments.label,
+                                                            label=self.label_name,
                                                             label_perfdata=self.label_perfdata,
-                                                            threshold=self.arguments.warning_threshold))
-                sys.exit(STATUS_WARNING)
+                                                            threshold=self.warning_threshold))
+                return STATUS_WARNING
             else:
                 print("OK: {value}{label} | {label_perfdata}={value};;;0".format(value=value,
                                                                                  label_perfdata=self.label_perfdata,
-                                                                                 label=self.arguments.label))
-                sys.exit(STATUS_OK)
+                                                                                 label=self.label_name))
+                return STATUS_OK
         except UnicodeDecodeError:
             print("UNKNOWN: can not read from sensor '{sensor_name}'! "
-                  "use {app_name} -h for further information!".format(sensor_name=self.arguments.sensor_name,
+                  "use {app_name} -h for further information!".format(sensor_name=self.sensor_name,
                                                                       app_name=os.path.basename(__file__)))
-            sys.exit(STATUS_UNKNOWN)
+            return STATUS_UNKNOWN
         except Exception as error:
             print("UNKNOWN: " + str(error))
-            sys.exit(STATUS_UNKNOWN)
+            return STATUS_UNKNOWN
+
+    def check_command(self):
+        return_code = self.check(sensor_name=self.sensor_name,
+                                 warning_threshold=self.warning_threshold,
+                                 critical_threshold=self.critical_threshold)
+        sys.exit(return_code)
 
 
 if __name__ == '__main__':
     check_pcmeasure = CheckPcMeasure()
-    check_pcmeasure.check()
+    check_pcmeasure.check_command()
diff --git a/tests.py b/tests.py
@@ -1,32 +1,23 @@
+#!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from collections import OrderedDict
+
 import socketserver
 import threading
-import time
+import unittest
+
 
 from unittest import TestCase
-from check_pcmeasure import (
-    MessPCCheck,
-    NAGIOS_OK, NAGIOS_WARNING, NAGIOS_CRITICAL, NAGIOS_UNKNOWN
-)
+from check_pcmeasure import CheckPcMeasure, STATUS_OK, STATUS_WARNING, STATUS_CRITICAL
 
 
 class AlwaysReturnTwentyOneDotSevenRequestHandler(socketserver.BaseRequestHandler):
     """ Handler which always returns 21.7 as value"""
 
     def handle(self):
         self.data = self.request.recv(2096).strip()
-        print(
-            "{client_ip} wrote: {data}".format(
-                client_ip=self.client_address[0],
-                data=self.data
-            )
-        )
-        self.request.sendall(
-            'valid=1;value= {expected_value};\n'.format(
-                expected_value=21.7
-            ).encode()
-        )
+        print("{client_ip} wrote: {data}".format(client_ip=self.client_address[0],
+                                                 data=self.data))
+        self.request.sendall('valid=1;value= {expected_value};\n'.format(expected_value=21.7).encode())
 
 
 class TestServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
@@ -46,51 +37,49 @@ def start_server(self):
 
     def setUp(self):
         self.start_server()
-        self.mess_pc_check = MessPCCheck(host=self.host, port=self.port)
+        self.mess_pc_check = CheckPcMeasure(host=self.host, port=self.port)
 
     def tearDown(self):
         self.stop_server()
+        self.mess_pc_check.close()
         del self.mess_pc_check
 
     def stop_server(self):
         self.server.shutdown()
         self.server.server_close()
 
     def test_01_get_value(self):
-        value = self.mess_pc_check._get_value(b"com1.1")
-        self.assertEqual(
-            value,
-            21.7
-        )
+        value = self.mess_pc_check._get_value("com1.1")
+        self.assertEqual(value, 21.7)
 
     def threshold_test(self, parameters):
-        result = self.mess_pc_check.check(
-            sensor_name=b"com1.1",
-            warning_threshold=parameters["warning"],
-            critical_threshold=parameters["critical"],
-        )
-        self.assertEqual(
-            result,
-            parameters["expected"]
-        )
+        try:
+            label = parameters["label"]
+        except KeyError:
+            label = False
+        result = self.mess_pc_check.check(sensor_name="com1.1",
+                                          warning_threshold=parameters["warning"],
+                                          critical_threshold=parameters["critical"],
+                                          label=label)
+        self.assertEqual(result, parameters["expected"])
 
     def test_02_check_no_thresholds_exceeded(self):
-        self.threshold_test({
-            "warning": 22.0,
-            "critical": 23.0,
-            "expected": NAGIOS_OK
-        })
+        self.threshold_test({"warning": 22.0, "critical": 23.0, "expected": STATUS_OK})
 
     def test_03_check_warning_threshold_exceeded(self):
-        self.threshold_test({
-            "warning": 21.5,
-            "critical": 23.0,
-            "expected": NAGIOS_WARNING
-        })
+        self.threshold_test({"warning": 21.5, "critical": 23.0, "expected": STATUS_WARNING})
 
     def test_04_check_critical_threshold_exceeded(self):
-        self.threshold_test({
-            "warning": 20.0,
-            "critical": 21.0,
-            "expected": NAGIOS_CRITICAL
-        })
+        self.threshold_test({"warning": 20.0, "critical": 21.0, "expected": STATUS_CRITICAL})
+
+    def test_05_check_no_thresholds_exceeded_with_label(self):
+        self.threshold_test({"warning": 22.0, "critical": 23.0, "expected": STATUS_OK, "label": "celsius"})
+
+    def test_06_check_warning_threshold_exceeded_with_label(self):
+        self.threshold_test({"warning": 21.5, "critical": 23.0, "expected": STATUS_WARNING, "label": "celsius"})
+
+    def test_07_check_critical_threshold_exceeded_with_label(self):
+        self.threshold_test({"warning": 20.0, "critical": 21.0, "expected": STATUS_CRITICAL, "label": "celsius"})
+
+if __name__ == '__main__':
+    unittest.main()
