Add asserts for the S extension level in tests

Guard these by checking the device SDK level first and adjusting the
expectations accordingly.

Note that the implementation is about to change, I'm just adding the
assert logic first so the actual impl changes are clearer in the
follow-up CL.

Bug: 173188089
Test: derive_sdk_test
Change-Id: Ic0fa1c06d2761c0390571900174810164b04e2cf
diff --git a/derive_sdk/derive_sdk_test.cpp b/derive_sdk/derive_sdk_test.cpp
index a0eb2a9..0f4ca72 100644
--- a/derive_sdk/derive_sdk_test.cpp
+++ b/derive_sdk/derive_sdk_test.cpp
@@ -21,6 +21,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/properties.h>
+#include <android-modules-utils/sdk_level.h>
 #include <gtest/gtest.h>
 #include <sys/stat.h>
 
@@ -47,21 +48,32 @@
     ASSERT_TRUE(android::base::WriteStringToFile(buf, path, true));
   }
 
+  void EXPECT_R(int n) {
+    int R = android::base::GetIntProperty("build.version.extensions.r", -1);
+    EXPECT_EQ(R, n);
+  }
+
+  void EXPECT_S(int n) {
+    int S = android::base::GetIntProperty("build.version.extensions.s", -1);
+    // Only expect the S extension level to be set on S+ devices.
+    EXPECT_EQ(S, android::modules::sdklevel::IsAtLeastS() ? n : -1);
+  }
+
   TemporaryDir dir_;
 };
 
-int R() {
-  return android::base::GetIntProperty("build.version.extensions.r", -1);
+TEST_F(DeriveSdkTest, CurrentSystemImageValue) {
+  EXPECT_R(0);
+  EXPECT_S(0);
 }
 
-TEST_F(DeriveSdkTest, CurrentSystemImageValue) { EXPECT_EQ(R(), 0); }
-
 TEST_F(DeriveSdkTest, OneApex) {
   MakeSdkVersion("a", 3);
 
   android::derivesdk::SetSdkLevels(dir());
 
-  EXPECT_EQ(R(), 3);
+  EXPECT_R(3);
+  EXPECT_S(3);
 }
 
 TEST_F(DeriveSdkTest, TwoApexes) {
@@ -70,7 +82,8 @@
 
   android::derivesdk::SetSdkLevels(dir());
 
-  EXPECT_EQ(R(), 3);
+  EXPECT_R(3);
+  EXPECT_S(3);
 }
 
 int main(int argc, char **argv) {