From a8cac44b677e440a552921794b51699edf4cf846 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=A9mie=20Galarneau?=
 <jeremie.galarneau@efficios.com>
Date: Thu, 15 Jun 2023 12:14:01 -0400
Subject: [PATCH] Tests: lttngtest: add methods to control session rotations
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jérémie Galarneau <jeremie.galarneau@efficios.com>
Change-Id: I8b8513fa37e6a51e034e5832499085c501e34ac4
---
 tests/utils/lttngtest/lttng.py    | 27 +++++++++++++++++++++++++++
 tests/utils/lttngtest/lttngctl.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tests/utils/lttngtest/lttng.py b/tests/utils/lttngtest/lttng.py
index 2b9dec5ad..6b86bbac5 100644
--- a/tests/utils/lttngtest/lttng.py
+++ b/tests/utils/lttngtest/lttng.py
@@ -292,6 +292,10 @@ class _Session(lttngctl.Session):
         # type: () -> None
         self._client._run_cmd("destroy '{session_name}'".format(session_name=self.name))
 
+    def rotate(self, wait=True):
+        # type: (bool) -> None
+        self._client.rotate_session_by_name(self.name, wait)
+
     @property
     def is_active(self):
         # type: () -> bool
@@ -480,6 +484,29 @@ class LTTngClient(logger._Logger, lttngctl.Controller):
         # type: () -> None
         self._run_cmd("destroy --all")
 
+    def rotate_session_by_name(self, name, wait=True):
+        self._run_cmd(
+            "rotate '{session_name}' {wait_option}".format(
+                session_name=name, wait_option="-n" if wait is False else ""
+            )
+        )
+
+    def schedule_size_based_rotation(self, name, size_bytes):
+        # type (str, int) -> None
+        self._run_cmd(
+            "enable-rotation --session '{session_name}' --size {size}".format(
+                session_name=name, size=size_bytes
+            )
+        )
+
+    def schedule_time_based_rotation(self, name, period_seconds):
+        # type (str, int) -> None
+        self._run_cmd(
+            "enable-rotation --session '{session_name}' --timer {period_seconds}s".format(
+                session_name=name, period_seconds=period_seconds
+            )
+        )
+
     @staticmethod
     def _mi_find_in_element(element, sub_element_name):
         # type: (xml.etree.ElementTree.Element, str) -> xml.etree.ElementTree.Element
diff --git a/tests/utils/lttngtest/lttngctl.py b/tests/utils/lttngtest/lttngctl.py
index c2afa885f..64c78b839 100644
--- a/tests/utils/lttngtest/lttngctl.py
+++ b/tests/utils/lttngtest/lttngctl.py
@@ -392,6 +392,11 @@ class Session(abc.ABC):
         # type: () -> bool
         pass
 
+    @abc.abstractmethod
+    def rotate(self):
+        # type: () -> None
+        raise NotImplementedError
+
     @abc.abstractproperty
     def kernel_pid_process_attribute_tracker(self):
         # type: () -> Type[ProcessIDProcessAttributeTracker]
@@ -543,3 +548,27 @@ class Controller(abc.ABC):
         List all sessions visible to the current user.
         """
         pass
+
+    @abc.abstractmethod
+    def rotate_session_by_name(self, name, wait=True):
+        # type: (str, bool) -> None
+        """
+        Rotate a session
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def schedule_size_based_rotation(self, name, size_bytes):
+        # type: (str, int) -> None
+        """
+        Schedule automatic size-based rotations.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def schedule_time_based_rotation(self, name, period_seconds):
+        # type: (str, int) -> None
+        """
+        Schedule automatic time-based rotations.
+        """
+        raise NotImplementedError
-- 
2.34.1