Added speedtest plugin, fixed database connection error.

collector.py

@@ -15,6 +15,7 @@ from plugins.system.memory_plugin import MemoryPlugin
 from plugins.system.network_plugin import NetworkPlugin
 from plugins.system.ping_plugin import PingPlugin
 from plugins.system.temperatures_plugin import TemperaturesPlugin
+from plugins.system.speedtest_plugin import SpeedtestPlugin
 
 
 class Collector(object):
@@ -29,6 +30,7 @@ class Collector(object):
             NetworkPlugin(),
             TemperaturesPlugin(),
             PingPlugin(),
+            SpeedtestPlugin(),
 
             # finance
             StocksPlugin(),
@@ -46,9 +48,10 @@ class Collector(object):
     def schedule_plugins(self):
         start_date = datetime.now() + timedelta(seconds=10)
         for plugin in self.plugins:
-            self.scheduler.add_job(plugin.execute, 'interval', 
+            self.scheduler.add_job(plugin.execute_wrapper, 'interval', 
                                    seconds=plugin.get_interval(),
-                                   start_date=start_date)
+                                   start_date=start_date,
+                                   name=plugin.__class__.__name__)
 
     def run(self):
         logging.basicConfig()
@@ -56,7 +59,8 @@ class Collector(object):
 
         models = self.collect_models()
         database.initialize_db()
-        database.DB.create_tables(models)
+        with database.DB.connection_context():
+            database.DB.create_tables(models)
 
         self.schedule_plugins()
         logging.info('Started.')

config.py

@@ -59,6 +59,9 @@ PING_HOSTS = [
     'tibich.com'
 ]
 
+### Speedtest
+SPEEDTEST_INTERVAL = 15 * 60        # every 15 min
+
 ### Stocks
 STOCKS_INTERVAL = 12 * 60 * 60      # updates daily
 

database.py

@@ -1,5 +1,6 @@
 from peewee import DatabaseProxy, Model
 from playhouse.db_url import connect
+from playhouse.pool import PooledPostgresqlExtDatabase
 import config
 
 DB = DatabaseProxy()

install.sh

@@ -33,6 +33,7 @@ After=network.target
 
 [Service]
 Type=simple
+User=$USER
 Restart=always
 RestartSec=1
 ExecStart=/usr/bin/env python3 $(pwd)/collector.py

plugins/plugin.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import Tuple
+import database
 
 
 class Plugin(ABC):
@@ -12,3 +13,7 @@ class Plugin(ABC):
     @abstractmethod
     def execute(self) -> None:
         pass
+
+    def execute_wrapper(self) -> None:
+        with database.DB.connection_context():
+            self.execute()

plugins/system/smart_plugin.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+
+import psutil
+from peewee import *
+from playhouse.shortcuts import model_to_dict
+
+import config
+from database import BaseModel
+from plugins.plugin import Plugin
+
+
+class SMART(BaseModel):
+    time = DateTimeField(index=True, default=datetime.utcnow)
+    drive = TextField(null=False)
+    attribute_id = IntegerField(null=False)
+    attribute_name = TextField(null=False)
+    value = IntegerField(null=False)
+    worst = IntegerField(null=False)
+    threshold = IntegerField(null=False)
+    raw = IntegerField(null=False)
+
+
+class SMARTPlugin(Plugin):
+    models = [SMART]
+
+    def get_interval(self):
+        return config.DISK_USAGE_INTERVAL
+
+    def execute(self):
+        for partition in psutil.disk_partitions():
+            usage = psutil.disk_usage(partition.mountpoint)
+            
+            entry = DiskUsage()
+            entry.partition = partition.device
+            entry.mountpoint = partition.mountpoint
+            entry.total = usage.total
+            entry.used = usage.used
+            entry.free = usage.free
+            entry.save()

plugins/system/speedtest_plugin.py

@@ -0,0 +1,57 @@
+import subprocess
+import re
+import subprocess
+from datetime import datetime
+
+import psutil
+from peewee import *
+from playhouse.shortcuts import model_to_dict
+
+import config
+from database import BaseModel
+from plugins.plugin import Plugin
+
+import json
+
+import logging
+
+class Speedtest(BaseModel):
+    time = DateTimeField(index=True, default=datetime.utcnow)
+    upload = IntegerField(null=False)
+    download = IntegerField(null=False)
+    latency = FloatField(null=True)
+    jitter = FloatField(null=True)
+    packetLoss = IntegerField(null=True)
+
+
+class SpeedtestPlugin(Plugin):
+    models = [Speedtest]
+
+    def get_interval(self):
+        return config.SPEEDTEST_INTERVAL
+
+    def execute(self):
+        command = ['/usr/local/bin/SpeedTest', '--output', 'json']
+        proc = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout = proc.stdout.decode()
+        stderr = proc.stderr.decode()
+
+        entry = Speedtest()
+        entry.download = 0
+        entry.upload = 0
+
+        if proc.returncode == 0:
+            try:
+                result = json.loads(stdout)
+                entry.download = int(float(result['download']))
+                entry.upload = int(float(result['upload']))
+                entry.latency = float(result['ping'])
+                entry.jitter = float(result['jitter'])
+            except BaseException as e:
+                logging.error(f"SpeedTest failed: {e}")
+
+        else:
+            logging.error(f"SpeedTest nonzero return: {proc.returncode}\n-----\n{stdout}\n{stderr}\n\n")
+            
+
+        entry.save()

requirements.txt

@@ -7,9 +7,12 @@ psycopg2-binary
 # Used in most system plugins
 psutil
 
+# Used by SMART plugin, requires 'smartmontools' installed + root access
+pySMART
+
 # robor plugin
 requests
 lxml
 
 # stocks plugin
-yfinance
+yfinance