From 1102e428568d1f8ef57a775f4d1c8a2bbe78a7e5 Mon Sep 17 00:00:00 2001
From: TheBurchLog <5104941+TheBurchLog@users.noreply.github.com>
Date: Thu, 11 Jul 2024 13:25:31 -0400
Subject: [PATCH] Add max history length

---
 brewtils/models.py  |  5 ++++-
 test/models_test.py | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/brewtils/models.py b/brewtils/models.py
index f9a1ce6f..5c2aa1f6 100644
--- a/brewtils/models.py
+++ b/brewtils/models.py
@@ -452,11 +452,14 @@ def __init__(self, heartbeat=None, history=None):
         self.heartbeat = heartbeat
         self.history = history or []
 
-    def set_status_heartbeat(self, status):
+    def set_status_heartbeat(self, status, max_history=None):
 
         self.heartbeat = datetime.utcnow()
         self.history.append(StatusHistory(status=copy.deepcopy(status), heartbeat=self.heartbeat))
 
+        if max_history and len(self.history) > max_history:
+            self.history = self.history[(max_history * -1):]
+
     def __str__(self):
         return self.heartbeat
 
diff --git a/test/models_test.py b/test/models_test.py
index ec3f40dc..0946f559 100644
--- a/test/models_test.py
+++ b/test/models_test.py
@@ -20,6 +20,7 @@
     RequestTemplate,
     LegacyRole,
     Subscriber,
+    StatusInfo,
     Topic,
 )
 from pytest_lazyfixture import lazy_fixture
@@ -713,3 +714,24 @@ def test_repr(self, topic1, subscriber1):
             topic1.name,
             [subscriber1],
         )
+
+
+class TestStatusInfo:
+
+    def test_max_history(self):
+        status_info = StatusInfo()
+
+        max_length = 5
+
+        for _ in range(10):
+            status_info.set_status_heartbeat("RUNNING", max_history=max_length)
+
+        assert len(status_info.history) == max_length
+
+    def test_history(self):
+        status_info = StatusInfo()
+
+        for _ in range(10):
+            status_info.set_status_heartbeat("RUNNING")
+
+        assert len(status_info.history) == 10