diff --git a/keymaster/3.0/vts/functional/keymaster_hidl_hal_test.cpp b/keymaster/3.0/vts/functional/keymaster_hidl_hal_test.cpp
index 12bfac5..54c5ad5 100644
--- a/keymaster/3.0/vts/functional/keymaster_hidl_hal_test.cpp
+++ b/keymaster/3.0/vts/functional/keymaster_hidl_hal_test.cpp
@@ -925,11 +925,7 @@
                                        &att_tee_enforced,                //
                                        &att_unique_id));
 
-    if (att_keymaster_version == 3) {
-        EXPECT_EQ(2U, att_attestation_version);
-    } else {
-        EXPECT_EQ(1U, att_attestation_version);
-    }
+    EXPECT_TRUE(att_attestation_version == 1 || att_attestation_version == 2);
 
     expected_sw_enforced.push_back(TAG_ATTESTATION_APPLICATION_ID,
                                    HidlBuf(app_id));