From 2bec959544872b3c61859a5fe3d52c44fb1b0f1d Mon Sep 17 00:00:00 2001
From: Katarzyna Treder <katarzyna.treder@h-partners.com>
Date: Mon, 25 Nov 2024 09:09:44 +0100
Subject: [PATCH] Add method for updating initramfs accordingly to OS version

Signed-off-by: Katarzyna Treder <katarzyna.treder@h-partners.com>
---
 test_tools/initramfs.py | 19 +++++++++++++++++++
 test_utils/os_utils.py  | 23 +++++++++++++++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 test_tools/initramfs.py

diff --git a/test_tools/initramfs.py b/test_tools/initramfs.py
new file mode 100644
index 0000000..2a4f259
--- /dev/null
+++ b/test_tools/initramfs.py
@@ -0,0 +1,19 @@
+#
+# Copyright(c) 2024 Huawei Technologies Co., Ltd.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+
+from core.test_run import TestRun
+from test_utils.os_utils import get_distro, Distro
+
+
+def update():
+    distro = get_distro()
+    TestRun.LOGGER.info("Updating initramfs")
+    match distro:
+        case Distro.DEBIAN | Distro.UBUNTU:
+            TestRun.executor.run_expect_success("update-initramfs -u")
+        case Distro.OPENEULER | Distro.CENTOS | Distro.REDHAT:
+            TestRun.executor.run_expect_success(
+                "dracut -f /boot/initramfs-$(uname -r).img $(uname -r)"
+            )
diff --git a/test_utils/os_utils.py b/test_utils/os_utils.py
index e9c01f1..cc7abd5 100644
--- a/test_utils/os_utils.py
+++ b/test_utils/os_utils.py
@@ -8,9 +8,9 @@ import math
 import posixpath
 import re
 import time
-from datetime import timedelta, datetime
 
-from aenum import IntFlag, Enum, IntEnum
+from datetime import timedelta, datetime
+from enum import IntFlag, Enum, IntEnum, StrEnum
 from packaging import version
 from typing import List
 
@@ -28,6 +28,14 @@ DEBUGFS_MOUNT_POINT = "/sys/kernel/debug"
 MEMORY_MOUNT_POINT = "/mnt/memspace"
 
 
+class Distro(StrEnum):
+    UBUNTU = "ubuntu"
+    DEBIAN = "debian"
+    REDHAT = "rhel"
+    OPENEULER = "openeuler"
+    CENTOS = "centos"
+
+
 class DropCachesMode(IntFlag):
     PAGECACHE = 1
     SLAB = 2
@@ -102,6 +110,17 @@ class SystemManagerType(Enum):
     systemd = 1
 
 
+def get_distro():
+    output = TestRun.executor.run(
+        "cat /etc/os-release | grep -e \"^ID=\" | awk -F= '{print$2}' | tr -d '\"'"
+    ).stdout.lower()
+
+    try:
+        return Distro(output)
+    except ValueError:
+        raise ValueError(f"Could not resolve distro name. Command output: {output}")
+
+
 def get_system_manager():
     output = TestRun.executor.run_expect_success("ps -p 1").stdout
     type = output.split('\n')[1].split()[3]