From ecb20c2da2f287a285d5725382e1f389e8fec35b Mon Sep 17 00:00:00 2001
From: Madhu Kanoor <mkanoor@redhat.com>
Date: Thu, 3 Aug 2023 14:08:05 -0400
Subject: [PATCH] fix: send session_stats if heartbeat enabled after a rule is
 triggered (#563)

[AAP-12987](https://issues.redhat.com/browse/AAP-12987)
Previously we were only sending the session stats on timed intervals and
not when a rule gets triggered. This was causing the lag as reported in
the ticket. With this PR we are now sending the session stats when a
rule fires in addition to the timed intervals.
---
 ansible_rulebook/rule_set_runner.py |  6 ++++++
 tests/e2e/test_websocket.py         | 12 ++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/ansible_rulebook/rule_set_runner.py b/ansible_rulebook/rule_set_runner.py
index ecf7a28a..91fa8c2f 100644
--- a/ansible_rulebook/rule_set_runner.py
+++ b/ansible_rulebook/rule_set_runner.py
@@ -26,6 +26,7 @@
     MessageNotHandledException,
     MessageObservedException,
 )
+from drools.ruleset import session_stats
 
 from ansible_rulebook.builtin import actions as builtin_actions
 from ansible_rulebook.conf import settings
@@ -244,6 +245,11 @@ async def _drain_actionplan_queue(self):
                 queue_item = await self.ruleset_queue_plan.plan.queue.get()
                 rule_run_at = run_at()
                 action_item = cast(ActionContext, queue_item)
+                if self.parsed_args and self.parsed_args.heartbeat > 0:
+                    await send_session_stats(
+                        self.event_log,
+                        session_stats(self.ruleset_queue_plan.ruleset.name),
+                    )
                 if len(action_item.actions) > 1:
                     task = asyncio.create_task(
                         self._run_multiple_actions(action_item, rule_run_at)
diff --git a/tests/e2e/test_websocket.py b/tests/e2e/test_websocket.py
index 5a40ce9b..f64b9b6f 100644
--- a/tests/e2e/test_websocket.py
+++ b/tests/e2e/test_websocket.py
@@ -60,7 +60,7 @@ async def test_websocket_messages():
     job_counter = 0
     action_counter = 0
     session_stats_counter = 0
-
+    stats = None
     while not queue.empty():
         data = await queue.get()
         assert data["path"] == endpoint
@@ -101,11 +101,11 @@ async def test_websocket_messages():
             assert stats["numberOfRules"] == 1
             assert stats["numberOfDisabledRules"] == 0
             assert data["activation_id"] == proc_id
-            if session_stats_counter == 2:
-                assert stats["rulesTriggered"] == 1
-                assert stats["eventsProcessed"] == 2000
-                assert stats["eventsMatched"] == 1
-                assert stats["eventsSuppressed"] == 1999
+
+    assert stats["rulesTriggered"] == 1
+    assert stats["eventsProcessed"] == 2000
+    assert stats["eventsMatched"] == 1
+    assert stats["eventsSuppressed"] == 1999
 
     assert ansible_event_counter == 9
     assert session_stats_counter >= 2