sysmon: Add sysmon.main tests

BUG=chromium:736911
TEST=None

Change-Id: If7d9bdeb08fb0fe0d231a940e44c57329945e861
Reviewed-on: https://chromium-review.googlesource.com/570633
Commit-Ready: Allen Li <ayatane@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
diff --git a/scripts/sysmon/mainlib.py b/scripts/sysmon/mainlib.py
new file mode 100644
index 0000000..8a3814f
--- /dev/null
+++ b/scripts/sysmon/mainlib.py
@@ -0,0 +1,116 @@
+# Copyright 2016 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Send system monitoring data to the timeseries monitoring API."""
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+import random
+import time
+
+from chromite.lib import commandline
+from chromite.lib import cros_logging as logging
+from chromite.lib import metrics
+from chromite.lib import ts_mon_config
+from infra_libs.ts_mon.common import interface
+
+from chromite.scripts.sysmon import git_metrics
+from chromite.scripts.sysmon import loop
+from chromite.scripts.sysmon import net_metrics
+from chromite.scripts.sysmon import osinfo_metrics
+from chromite.scripts.sysmon import proc_metrics
+from chromite.scripts.sysmon import prod_metrics
+from chromite.scripts.sysmon import puppet_metrics
+from chromite.scripts.sysmon import system_metrics
+
+logger = logging.getLogger(__name__)
+
+
+class _MetricCollector(object):
+  """Metric collector class."""
+
+  def __init__(self, collect_prod_hosts=False):
+    self._collect_osinfo = _TimedCallback(
+        callback=osinfo_metrics.collect_os_info,
+        interval=60 * 60)
+    if collect_prod_hosts:
+      logger.info(u'Enabling prod host metric collection.')
+      self._collect_prod_hosts = _TimedCallback(
+          callback=prod_metrics.collect_prod_hosts,
+          interval=10 * 60)
+    else:
+      self._collect_prod_hosts = lambda: None
+
+  def __call__(self):
+    """Collect metrics."""
+    system_metrics.collect_uptime()
+    system_metrics.collect_cpu_info()
+    system_metrics.collect_disk_info()
+    system_metrics.collect_mem_info()
+    net_metrics.collect_net_info()
+    proc_metrics.collect_proc_info()
+    system_metrics.collect_load_avg()
+    puppet_metrics.collect_puppet_summary()
+    git_metrics.collect_git_metrics()
+    self._collect_prod_hosts()
+    self._collect_osinfo()
+    system_metrics.collect_unix_time()  # must be just before flush
+    metrics.Flush()
+
+
+class _TimedCallback(object):
+  """Limits callback to one call in a given interval."""
+
+  def __init__(self, callback, interval):
+    """Initialize instance.
+
+    Args:
+      callback: function to call
+      interval: Number of seconds between allowed calls
+    """
+    self._callback = callback
+    self._interval = interval
+    self._last_called = time.time() - interval
+
+  def __call__(self):
+    if time.time() >= self._next_call:
+      self._callback()
+      self._last_called = time.time()
+
+  @property
+  def _next_call(self):
+    return self._last_called + self._interval
+
+
+def main():
+  parser = commandline.ArgumentParser(
+      description=__doc__,
+      default_log_level='DEBUG')
+  parser.add_argument(
+      '--interval',
+      default=60,
+      type=int,
+      help='time (in seconds) between sampling system metrics')
+  parser.add_argument(
+      '--collect-prod-hosts',
+      action='store_true',
+      help='Enable collection of prod host metrics, like roles')
+  opts = parser.parse_args()
+  opts.Freeze()
+
+  # Wait a random amount of time before starting the loop in case sysmon
+  # is started at exactly the same time on all machines.
+  time.sleep(random.uniform(0, opts.interval))
+
+  # This call returns a context manager that doesn't do anything, so we
+  # ignore the return value.
+  ts_mon_config.SetupTsMonGlobalState('sysmon', auto_flush=False)
+  # The default prefix is '/chrome/infra/'.
+  interface.state.metric_name_prefix = (interface.state.metric_name_prefix
+                                        + 'chromeos/sysmon/')
+
+  collector = _MetricCollector(collect_prod_hosts=opts.collect_prod_hosts)
+  loop.SleepLoop(callback=collector,
+                 interval=opts.interval).loop_forever()