From 4e8575cbe1b915613f607f4d5803f8aa87972b71 Mon Sep 17 00:00:00 2001
From: sqozz <sqozz@geekify.de>
Date: Fri, 11 May 2018 20:12:42 +0200
Subject: [PATCH] Add mass metric

---
 schlafana.py | 40 ++++++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/schlafana.py b/schlafana.py
index bbf47c6..f433241 100644
--- a/schlafana.py
+++ b/schlafana.py
@@ -5,7 +5,7 @@ import logging
 from pprint import pprint
 from websocket import create_connection
 from prometheus_client import Counter, Gauge, Summary, start_http_server
-from prometheus_client.core import CounterMetricFamily, REGISTRY
+from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY
 import time
 
 MESSAGES = { "GameInfo" : 0x00,
@@ -44,6 +44,16 @@ class BotRevisionCountCollector(object):
         yield c
 
 
+class BotMassGaugeCollector(object):
+    def collect(self):
+        c = GaugeMetricFamily("mass", "Mass of each snake", labels=["name", "snake_id", "code_id"])
+        for user in game.users.users:
+            user = game.users.users.get(user)
+            rev = user.revs.get(user.current_rev)
+            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.mass)
+        yield c
+
+
 def main():
     global game
     game = Game()
@@ -61,6 +71,7 @@ def main():
     REGISTRY.register(BotDeathCollector())
     REGISTRY.register(BotKillCollector())
     REGISTRY.register(BotRevisionCountCollector())
+    REGISTRY.register(BotMassGaugeCollector())
     start_http_server(9000)
     logger.info("Start prometheus scrape endpoint")
 
@@ -177,14 +188,7 @@ class BotUser():
 
     def update(self, data):
         rev_id = data.get("db_id")
-        worm_id = data.get("id") #TODO: does this work?
-        if "data" in data.keys():
-            user = game.users.get_user_by_worm(worm_id)
-            current_rev = user.revs.get(user.current_rev)
-            current_rev.carrion_food_consumed = data.get("data").get("c")
-            current_rev.natural_food_consumed = data.get("data").get("n")
-            current_rev.hunted_food_consumed = data.get("data").get("h")
-            pdb.set_trace()
+        worm_id = data.get("id")
         if rev_id in self.revs.keys(): # revision already known
             rev = self.revs.get(rev_id)
             rev.update(data)
@@ -212,13 +216,17 @@ class BotRevision():
         pass
 
     def update(self, data):
-        if "food" in data.keys():
-            pdb.set_trace()
-        self.current_worm = data.get("id")
-        self.rev_id = data.get("db_id")
-        self.mass = data.get("mass")
-        self.worms.update({self.current_worm : [self.rev_id]})
-        pass
+        if "data" in data.keys():
+            self.mass = data.get("data").get("m")
+            self.carrion_food_consumed = data.get("data").get("c")
+            self.natural_food_consumed = data.get("data").get("n")
+            self.hunted_food_consumed = data.get("data").get("h")
+        else:
+            self.current_worm = data.get("id")
+            self.rev_id = data.get("db_id")
+            self.mass = data.get("mass")
+            self.worms.update({self.current_worm : [self.rev_id]})
+            pass
 
 
 class UserList():