From f66b9c4556847964298afe22cb149a22c64e8855 Mon Sep 17 00:00:00 2001
From: Nemo <me@captnemo.in>
Date: Sun, 1 Jul 2018 03:06:22 +0530
Subject: [PATCH] Switches to new prometheus formatter

---
 index.js  |  9 ++++++---
 prom.js   | 30 ++++++++++++++++++++++++++++++
 server.js | 31 ++++---------------------------
 3 files changed, 40 insertions(+), 30 deletions(-)
 create mode 100644 prom.js

diff --git a/index.js b/index.js
index 3bdd505..804540e 100644
--- a/index.js
+++ b/index.js
@@ -6,6 +6,7 @@ const MY_PACKAGE_SELECTOR_ID =
 const DATA_SELECTOR = 'packagecol3';
 
 const DATA_USAGE_REGEX = /\d+\.\d{0,2}/g;
+const KEYS = ['live', 'flexibytes'];
 var browser;
 
 async function getUsage() {
@@ -28,18 +29,20 @@ async function getUsage() {
     await page.waitFor(3000),
       (dataUsage = await page.evaluate(sel => {
         let elements = document.getElementsByClassName(sel);
+
+        if (!elements || elements.length < 4) {
+          return '0.00 0.00';
+        }
         let usage = {
           live: elements[3].innerText,
           aggregate: '0.00 GB (Quota 800.00 GB)',
         };
-        if (elements.length >= 5) {
+        if (elements.length >= 6) {
           usage['flexibytes'] = elements[5].innerText;
         }
         return usage;
       }, DATA_SELECTOR));
 
-    const KEYS = ['live', 'flexibytes'];
-
     // ['0.00 GB (Quota 800.00 GB)', '102.58 GB(Quota 100.00 GB)']
     KEYS.map(key => {
       dataUsage[key] = dataUsage[key]
diff --git a/prom.js b/prom.js
new file mode 100644
index 0000000..e7f5b4c
--- /dev/null
+++ b/prom.js
@@ -0,0 +1,30 @@
+const pClient = require('prom-client');
+const KEYS = ['live', 'flexibytes', 'aggregate'];
+
+// Initial Setup
+let gauges = {};
+pClient.collectDefaultMetrics({ timeout: 60000 });
+
+KEYS.map(key => {
+	gauges[`${key}_used`] = new pClient.Gauge({
+		name: `act_fup_${key}_usage_bytes`,
+		help: `ACT ${key} usage in bytes (precision GB)`,
+	});
+
+	gauges[`${key}_total`] = new pClient.Gauge({
+		name: `act_fup_${key}_total_usage_bytes`,
+		help: `ACT ${key} usage in bytes (precision GB)`,
+	});
+});
+
+module.exports = {
+	format: function(metrics) {
+		KEYS.map(key => {
+			gauges[`${key}_used`].set(metrics[key].usedBytes);
+			gauges[`${key}_total`].set(metrics[key].totalBytes);
+		});
+
+		return pClient.register.metrics();
+	},
+	contentType: pClient.register.contentType,
+};
diff --git a/server.js b/server.js
index 29d755d..d80dc21 100644
--- a/server.js
+++ b/server.js
@@ -1,38 +1,15 @@
 const http = require('http');
 const port = 3000;
-const pClient = require('prom-client');
 const metrics = require('./index');
-
-pClient.collectDefaultMetrics({ timeout: 60000 });
-
-let usedGauge = new pClient.Gauge({
-    name: 'act_fup_usage_bytes',
-    help: 'ACT current usage in bytes (precision GB)',
-});
-
-let totalGauge = new pClient.Gauge({
-    name: 'act_fup_max_bytes',
-    help: 'ACT FUP limit in bytes (precision GB)',
-});
+const promFormatter = require('./prom');
 
 const requestHandler = async (req, res) => {
     let date = new Date(Date.now()).toLocaleString();
     console.log(`${date}: ${req.url}`);
     switch (req.url) {
         case '/metrics':
-            let m = await metrics.getUsage();
-            // TODO: Switch to the correct err, res pattern with promise
-            if (m !== null) {
-                usedGauge.set(m.usedBytes);
-                totalGauge.set(m.totalBytes);
-
-                res.setHeader('Content-Type', pClient.register.contentType);
-                res.end(pClient.register.metrics());
-            } else {
-                res.sendStatus(500);
-                res.end('Scrape failed');
-            }
-
+            res.setHeader('Content-Type', promFormatter.contentType);
+            res.end(promFormatter.format(await metrics.getUsage()));
             break;
         default:
             res.writeHead(302, {
@@ -48,7 +25,7 @@ const server = http.createServer(requestHandler);
 
 server.listen(port, err => {
     if (err) {
-        return console.log('something bad happened', err);
+        return console.log('could not initialize web server', err);
     }
 
     console.log(`server is listening on ${port}`);