check_computer_memory: extend API

Add the possibility to choose unit between byte or percent,
and to choose between checking used or free memory.

slapos/promise/check_computer_memory/__init__.py

 #!/usr/bin/env python
 
 """
-Check if memory usage is lower than 80% of total memory.
+Check if memory usage is lower a given threshold.
 
 Uses:
 - /proc/meminfo
@@ -24,7 +24,7 @@ def getMemoryInfo(database, time, date):
     result = zip(*query_result)
     if not result or not result[0][0]:
       return (None, "couldn't fetch total memory, collectordb is empty?")
-    memory_info['total'] = result[0][0]
+    memory_info['total'] = int(result[0][0])  # in byte
 
     # fetch free and used memory 
     where_query = "time between '%s:00' and '%s:30' " % (time, time)
@@ -32,18 +32,50 @@ def getMemoryInfo(database, time, date):
     result = zip(*query_result)
     if not result or not result[0][0]: 
       return (None, "couldn't fetch free memory")
-    memory_info['free'] = result[0][0]
+    memory_info['free'] = int(result[0][0])  # in byte
     if not result or not result[1][0]: 
       return (None, "couldn't fetch used memory")
-    memory_info['used'] = result[1][0]
+    memory_info['used'] = int(result[1][0])  # in byte
   finally:
     database.close()
 
+  memory_info["used_percent"] = memory_info["used"] * 100.0 / memory_info["total"]
+  memory_info["free_percent"] = memory_info["free"] * 100.0 / memory_info["total"]
   return (memory_info, "")
 
+def checkMemoryUsage(database_path, time, date, threshold, key="used", unit="byte"):
+  if key not in ("used", "free"):
+    raise ValueError("invalid key")
+  if unit not in ("byte", "percent"):
+    raise ValueError("invalid unit")
+  memory_info, error = getMemoryInfo(database_path, time, date)
+  if error:
+    return (False, "error", error)
+  if unit == "byte":
+    if key == "used":
+      if memory_info["used"] <= threshold:
+        return (True, "OK - memory used: {used}B of {total}B ({used_percent:.1f}%)".format(**memory_info), "")
+      return (False, "Low available memory - usage: {used}B of {total}B ({used_percent:.1f}%)".format(**memory_info), "")
+    else:  # if key == "free":
+      if memory_info["free"] > threshold:
+        return (True, "OK - free memory: {free}B of {total}B ({free_percent:.1f}%)".format(**memory_info), "")
+      return (False, "Low available memory - free: {free}B of {total}B ({free_percent:.1f}%)".format(**memory_info), "")
+  else:  # if unit == "percent":
+    if key == "used":
+      if memory_info["used_percent"] <= threshold:
+        return (True, "OK - memory used: {used_percent:.1f}% ({used}B of {total}B)".format(**memory_info), "")
+      return (False, "Low available memory - usage: {used_percent:.1f}% ({used}B of {total}B)".format(**memory_info), "")
+    else:  # if key == "free":
+      if memory_info["free_percent"] > threshold:
+        return (True, "OK - free memory: {free_percent:.1f}% ({free}B of {total}B)".format(**memory_info), "")
+      return (False, "Low available memory - free: {free_percent:.1f}% ({free}B of {total}B)".format(**memory_info), "")
+
 def main():
   parser = argparse.ArgumentParser()
   parser.add_argument("-db", "--collectordb", required=True)
+  parser.add_argument("--threshold", required=True, type=float)
+  parser.add_argument("--key", default="used", choices=("used", "free"))
+  parser.add_argument("--unit", default="byte", choices=("byte", "percent"))
   args = parser.parse_args()
 
   # get last minute
@@ -55,17 +87,18 @@ def main():
   db_path = args.collectordb
 
   if db_path.endswith("collector.db"):db_path=db_path[:-len("collector.db")]
-  memory, error = getMemoryInfo(db_path, currenttime, currentdate)
+
+  result, message, error = checkMemoryUsage(
+    db_path, currenttime, currentdate,
+    threshold=args.threshold,
+    key=args.key,
+    unit=args.unit,
+  )
   if error:
     print error
     return 0
-  threshold = float(memory['total']) * 0.2
-
-  if memory['used'] <= threshold:
-    print "OK - memory used: {:d}% ({}B of {}B)".format(memory["used"] * 100 / memory["total"], memory["used"], memory["total"])
-    return 0
-  print "Low memory - usage: {:d}% ({}B of {}B)".format(memory["used"] * 100 / memory["total"], memory["used"], memory["total"])
-  return 1
+  print message
+  return 1 if result else 0
 
 if __name__ == "__main__":
   sys.exit(main())

slapos/test/promise/test_check_computer_memory.py

@@ -30,7 +30,7 @@ import os
 import sqlite3
 
 from slapos.test.promise import data
-from slapos.promise.check_computer_memory import getMemoryInfo
+from slapos.promise.check_computer_memory import getMemoryInfo, checkMemoryUsage
 
 total_memory_fetch_failure_message = "couldn't fetch total memory, collectordb is empty?"
 
@@ -49,13 +49,52 @@ class TestComputerMemory(unittest.TestCase):
     self.conn.close() 
 
   def test_check_memory(self):
-    self.assertEquals(({'total': 33705312256.0, 'used': 33139023872.0, 'free': 566288384.0}, ""),
+    self.assertEquals(
+      ({
+        'total': 33705312256,
+        'used': 33139023872,
+        'used_percent': 98.31988388151132,
+        'free': 566288384,
+        'free_percent': 1.6801161184886904,
+      }, ""),
       getMemoryInfo('/tmp', '00:02', '2017-09-15'))
+    self.assertEquals(
+      (True, "OK - memory used: 33139023872B of 33705312256B (98.3%)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=33500000000),
+    )
+    self.assertEquals(
+      (True, "OK - memory used: 98.3% (33139023872B of 33705312256B)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=99, unit="percent"),
+    )
+    self.assertEquals(
+      (True, "OK - free memory: 566288384B of 33705312256B (1.7%)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=500000000, key="free"),
+    )
+    self.assertEquals(
+      (True, "OK - free memory: 1.7% (566288384B of 33705312256B)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=1, key="free", unit="percent"),
+    )
+    self.assertEquals(
+      (False, "Low available memory - usage: 33139023872B of 33705312256B (98.3%)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=33000000000),
+    )
+    self.assertEquals(
+      (False, "Low available memory - usage: 98.3% (33139023872B of 33705312256B)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=98, unit="percent"),
+    )
+    self.assertEquals(
+      (False, "Low available memory - free: 566288384B of 33705312256B (1.7%)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=600000000, key="free"),
+    )
+    self.assertEquals(
+      (False, "Low available memory - free: 1.7% (566288384B of 33705312256B)", ""),
+      checkMemoryUsage('/tmp', '00:02', '2017-09-15', threshold=2, key="free", unit="percent"),
+    )
 
   def test_check_memory_with_unavailable_dates(self):
     self.assertEquals(
-      (None, total_memory_fetch_failure_message),
-      getMemoryInfo('/tmp', '18:00', '2017-09-14'),
+      (False, "error", total_memory_fetch_failure_message),
+      checkMemoryUsage('/tmp', '18:00', '2017-09-14', 1.0),
     )
 
   def tearDown(self):