From 5e87e135ae9067a29da9577f8eaeddddacb6a4bf Mon Sep 17 00:00:00 2001
From: Patrick Zhang <patrzhan@amazon.com>
Date: Tue, 7 Nov 2023 14:46:36 -0800
Subject: [PATCH] feat: add steps for ram and cpu use

---
 .../pom.xml                                   |  5 ++
 .../testing/features/SystemSteps.java         | 88 +++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 aws-greengrass-testing-features/aws-greengrass-testing-features-api/src/main/java/com/aws/greengrass/testing/features/SystemSteps.java

diff --git a/aws-greengrass-testing-features/aws-greengrass-testing-features-api/pom.xml b/aws-greengrass-testing-features/aws-greengrass-testing-features-api/pom.xml
index 87c09139..8cc2e7c9 100644
--- a/aws-greengrass-testing-features/aws-greengrass-testing-features-api/pom.xml
+++ b/aws-greengrass-testing-features/aws-greengrass-testing-features-api/pom.xml
@@ -200,6 +200,11 @@
             <artifactId>guice</artifactId>
             <version>${guice.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.github.oshi</groupId>
+            <artifactId>oshi-core</artifactId>
+            <version>6.4.7</version>
+        </dependency>
         <dependency>
             <groupId>org.immutables</groupId>
             <artifactId>value</artifactId>
diff --git a/aws-greengrass-testing-features/aws-greengrass-testing-features-api/src/main/java/com/aws/greengrass/testing/features/SystemSteps.java b/aws-greengrass-testing-features/aws-greengrass-testing-features-api/src/main/java/com/aws/greengrass/testing/features/SystemSteps.java
new file mode 100644
index 00000000..80a957f4
--- /dev/null
+++ b/aws-greengrass-testing-features/aws-greengrass-testing-features-api/src/main/java/com/aws/greengrass/testing/features/SystemSteps.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.aws.greengrass.testing.features;
+
+import io.cucumber.guice.ScenarioScoped;
+import io.cucumber.java.en.Then;
+import io.cucumber.java.en.When;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import oshi.SystemInfo;
+import oshi.hardware.CentralProcessor;
+import oshi.hardware.GlobalMemory;
+import oshi.hardware.HardwareAbstractionLayer;
+
+import java.util.ArrayList;
+
+
+@ScenarioScoped
+public class SystemSteps {
+    private static final Logger LOGGER = LogManager.getLogger(SystemSteps.class);
+    private ArrayList<Double> ramList = new ArrayList<Double>();
+    private ArrayList<Double> cpuList = new ArrayList<Double>();
+    private SystemInfo si = new SystemInfo();
+    private HardwareAbstractionLayer hal = si.getHardware();
+    private CentralProcessor cpu = hal.getProcessor();
+    private GlobalMemory ram = hal.getMemory();
+
+    /**
+     * Record the current CPU or RAM usage depending on the step.
+     *
+     * @param stat the statistic to record, must be CPU or RAM
+     * @throws IllegalArgumentException when parameter is not CPU or RAM
+     */
+    @When("I record the device's {word} usage statistic")
+    public void recordCpuOrRam(String stat) throws IllegalArgumentException {
+        switch (stat) {
+            case "CPU":
+                double cpuLoad = cpu.getSystemCpuLoad(500) * 100;
+                LOGGER.info("System CPU load recorded as: " + cpuLoad + "%");
+                cpuList.add(0, cpuLoad);
+                break;
+            case "RAM":
+                double usedRam = (ram.getTotal() - ram.getAvailable()) / (1024 * 1024);
+                LOGGER.info("Used RAM recorded as: " + usedRam + " MB");
+                ramList.add(0, usedRam);
+                break;
+            default:
+                throw new IllegalArgumentException("Please specify either CPU or RAM to track");
+        }
+    }
+
+    /**
+     * Check difference between last two records and verify it is below a threshold.
+     *
+     * @param stat the statistic to record, must be CPU or RAM
+     * @param threshold the threshold to assert the difference is under, % for CPU and MB for RAM
+     * @throws Exception when step fails due to exceeding the provided threshold or sample steps haven't run twice
+     * @throws IllegalArgumentException when stat param is not CPU or RAM
+     */
+    @Then("the difference in the last two {word} samples is less than {int}")
+    public void checkSampleDiff(String stat, int threshold) throws Exception, IllegalArgumentException {
+        switch (stat) {
+            case "CPU":
+                if (cpuList.size() < 2) {
+                    throw new Exception("Need at least two CPU samples first.");
+                }
+                double cpuDiff = cpuList.get(0) - cpuList.get(1);
+                if (cpuDiff >= threshold) {
+                    throw new Exception("CPU use was above the provided threshold.");
+                }
+                break;
+            case "RAM":
+                if (ramList.size() < 2) {
+                    throw new Exception("Need at least two RAM samples first.");
+                }
+                double ramDiff = ramList.get(0) - ramList.get(1);
+                if (ramDiff >= threshold) {
+                    throw new Exception("RAM use was above the provided threshold");
+                }
+                break;
+            default:
+                throw new IllegalArgumentException("Please specify either CPU or RAM to check");
+        }
+    }
+}