From 71843c94147661b3f454569b10a00a5a1b4845c7 Mon Sep 17 00:00:00 2001
From: Torben Hansen <50673096+torben-hansen@users.noreply.github.com>
Date: Wed, 4 Dec 2024 08:48:33 -0800
Subject: [PATCH] Fix tests

---
 crypto/fipsmodule/rand/new_rand_test.cc | 13 ++++++-------
 crypto/ube/ube_test.cc                  | 13 ++++++-------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/crypto/fipsmodule/rand/new_rand_test.cc b/crypto/fipsmodule/rand/new_rand_test.cc
index 57c6f5577b..cc66732e7b 100644
--- a/crypto/fipsmodule/rand/new_rand_test.cc
+++ b/crypto/fipsmodule/rand/new_rand_test.cc
@@ -21,8 +21,9 @@ class newRandTest : public ::testing::Test {
   public:
     void SetUp() override {
       uint64_t current_generation_number = 0;
-      ube_detection_supported_ = CRYPTO_get_ube_generation_number(
-                                  &current_generation_number);
+      if (CRYPTO_get_ube_generation_number(&current_generation_number) == 1) {
+        ube_detection_supported_ = true;
+      }
     }
 
     void TearDown() override {
@@ -34,10 +35,8 @@ class newRandTest : public ::testing::Test {
       return ube_detection_supported_;
     }
 
-    void allowMockedUbeIfNecessary(void) {
-      if (ube_detection_supported_) {
-        allow_mocked_ube_detection_FOR_TESTING();
-      }
+    void allowMockedUbe(void) {
+      allow_mocked_ube_detection_FOR_TESTING();
     }
 
     bool ube_detection_supported_ = false;
@@ -191,7 +190,7 @@ static void MockedUbeDetection(std::function<void(uint64_t)> set_detection_metho
 
 TEST_F(newRandTest, UbeDetectionMocked) {
 
-  allowMockedUbeIfNecessary();
+  allowMockedUbe();
 
   MockedUbeDetection(
     [](uint64_t gn) {
diff --git a/crypto/ube/ube_test.cc b/crypto/ube/ube_test.cc
index 3314e1779e..b425e659a8 100644
--- a/crypto/ube/ube_test.cc
+++ b/crypto/ube/ube_test.cc
@@ -10,8 +10,9 @@ class ubeTest : public ::testing::Test {
   public:
     void SetUp() override {
       uint64_t current_generation_number = 0;
-      ube_detection_supported_ = CRYPTO_get_ube_generation_number(
-                                  &current_generation_number);
+      if (CRYPTO_get_ube_generation_number(&current_generation_number) == 1) {
+        ube_detection_supported_ = true;
+      }
     }
 
     void TearDown() override {
@@ -23,10 +24,8 @@ class ubeTest : public ::testing::Test {
       return ube_detection_supported_;
     }
 
-    void allowMockedUbeIfNecessary(void) {
-      if (ube_detection_supported_) {
-        allow_mocked_ube_detection_FOR_TESTING();
-      }
+    void allowMockedUbe(void) {
+      allow_mocked_ube_detection_FOR_TESTING();
     }
 
     bool ube_detection_supported_ = false;
@@ -56,7 +55,7 @@ TEST_F(ubeTest, BasicTests) {
 
 TEST_F(ubeTest, MockedMethodTests) {
 
-  allowMockedUbeIfNecessary();
+  allowMockedUbe();
 
   uint64_t generation_number = 0;
   uint64_t cached_generation_number = 0;