add unit testing

Author: WGussev

Diff for: power_watcher/power_watcher.py

@@ -50,7 +50,7 @@ def update(self, value):
         self.value = value
 
 
-def _watch_power(logfile: Path = None, sender: Connection = None, display: bool = True):
+def _watch_power(logfile: Path = None, sender: Connection = None, display: bool = False):
     """
     Poll GPU and log/display current power consumption.
     Update frequency: every 1 second.
@@ -60,8 +60,6 @@ def _watch_power(logfile: Path = None, sender: Connection = None, display: bool
     :param display: display consumption in terminal.
     :return: None
     """
-    if (logfile is None) and (display is False):
-        raise ValueError('You should log and/or display consumption')
 
     total = 0
     killer = _GracefulKiller()
@@ -79,6 +77,7 @@ def _watch_power(logfile: Path = None, sender: Connection = None, display: bool
         if logfile is not None:
             f.write(f'{datetime.now()} {power}\n')
         time.sleep(1)
+    print(total)
     if display:
         print('', end='\n')
     if sender is not None:
@@ -97,10 +96,10 @@ class PowerWatcher:
     pw.total  # get results
     """
 
-    def __init__(self, logfile: Path = None, display: bool = True):
+    def __init__(self, logfile: Path = None, display: bool = False):
         """
         :param logfile: logfile path.
-        :param display: consumption display toggle.
+        :param display: display consumption in the terminal.
         """
         self.logfile = logfile
         self.display = display
@@ -130,4 +129,4 @@ def start(self):
     def stop(self):
         """Stop manually."""
         self.watcher.terminate()
-        self.total.update(self.recv_end.recv())
+        self.total.update(self.recv_end.recv())

Diff for: test.py

@@ -0,0 +1,44 @@
+import unittest
+from importlib import import_module
+from multiprocessing import Process, Pipe
+from time import sleep
+from unittest.mock import patch
+
+
+class FirstTest(unittest.TestCase):
+    pw = import_module('power_watcher')
+    pw = pw.power_watcher
+
+    def test_value_container(self):
+        container = self.pw._ValueContainer()
+        self.assertEqual(container.value, 0)
+        container.update(7)
+        self.assertEqual(container.value, 7)
+        container.update(0.45)
+        self.assertEqual(container.value, 0.45)
+
+    def test_process_killer(self):
+        def f(sender):
+            killer = self.pw._GracefulKiller()
+            while killer.kill_now is False:
+                pass
+            sender.send(killer)
+
+        recv_end, send_end = Pipe(False)
+        watcher = Process(target=f, args=(send_end,))
+        watcher.start()
+        sleep(1)
+        watcher.terminate()
+        kn = recv_end.recv()
+        self.assertEqual(kn.kill_now, True)
+
+    @patch('power_watcher.power_watcher.nvmlInit')
+    @patch('power_watcher.power_watcher.nvmlDeviceGetHandleByIndex')
+    @patch('power_watcher.power_watcher.nvmlDeviceGetPowerUsage', return_value=3600 * 1000 * 1000)
+    def test_context_manager(self, mock_init, mock_handle, mock_power):
+        recv_end, send_end = Pipe(False)
+        watcher = Process(target=self.pw._watch_power, args=(None, send_end, False))
+        watcher.start()
+        sleep(1)
+        watcher.terminate()
+        self.assertEqual(recv_end.recv(), 1.0)