chore(modelarmor): Added floor settings tests (#10106)

* chore(modelarmor): Added floorsettings tests and marked them as skip for now

* chore(modelarmor): refactored TODO messages

* chore(tests): enabled floor settings tests

* add default variables for model armor tests

* fix(test): update method of folder floor setting

---------

Co-authored-by: Katie McLaughlin <[email protected]>
Co-authored-by: Harsh Nasit <[email protected]>

Author: 3 people

.kokoro/tests/run_tests.sh

@@ -45,7 +45,7 @@ fi
 if [[ "$SCRIPT_DEBUG" != "true" ]]; then
     # Update `gcloud` and log versioning for debugging
     apt update && apt -y upgrade google-cloud-sdk
-    
+
     echo "********** GIT INFO ***********"
     git version
     echo "********** GCLOUD INFO ***********"
@@ -67,7 +67,10 @@ if [[ "$SCRIPT_DEBUG" != "true" ]]; then
     # For Cloud Run filesystem sample
     export FILESTORE_IP_ADDRESS=$(gcloud secrets versions access latest --secret fs-app)
     export MNT_DIR=$PWD/run/filesystem
-    
+    # For Model Armor tests
+    export MA_FOLDER_ID=695279264361
+    export MA_ORG_ID=951890214235
+
     SECRET_FILES=("java-docs-samples-service-account.json" \
     "java-aiplatform-samples-secrets.txt" \
     "java-automl-samples-secrets.txt" \
@@ -88,7 +91,7 @@ if [[ "$SCRIPT_DEBUG" != "true" ]]; then
 
     # create secret dir
     mkdir -p "${KOKORO_GFILE_DIR}/secrets"
-    
+
     for SECRET in "${SECRET_FILES[@]}"; do
       # grab latest version of secret
       gcloud secrets versions access latest --secret="${SECRET%.*}" > "${KOKORO_GFILE_DIR}/secrets/$SECRET"
@@ -166,7 +169,7 @@ test_prog="$PWD/.kokoro/tests/run_test_java.sh"
 
 git config --global --add safe.directory $PWD
 
-# Use btlr to run all the tests in each folder 
+# Use btlr to run all the tests in each folder
 echo "btlr" "${btlr_args[@]}" -- "${test_prog}"
 btlr "${btlr_args[@]}" -- "${test_prog}"
 

modelarmor/pom.xml

@@ -43,7 +43,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>libraries-bom</artifactId>
-        <version>26.60.0</version>
+        <version>26.64.0</version>
         <type>pom</type>
         <scope>import</scope>
       </dependency>

modelarmor/src/main/java/modelarmor/GetFolderFloorSetting.java

@@ -38,7 +38,7 @@ public static FloorSetting getFolderFloorSetting(String folderId) throws IOExcep
     // Initialize client that will be used to send requests. This client only
     // needs to be created once, and can be reused for multiple requests.
     try (ModelArmorClient client = ModelArmorClient.create()) {
-      String name = FloorSettingName.of(folderId, "global").toString();
+      String name = FloorSettingName.ofFolderLocationName(folderId, "global").toString();
 
       GetFloorSettingRequest request = GetFloorSettingRequest.newBuilder().setName(name).build();
 

modelarmor/src/test/java/modelarmor/SnippetsIT.java

@@ -28,6 +28,8 @@
 import com.google.cloud.modelarmor.v1.FilterConfig;
 import com.google.cloud.modelarmor.v1.FilterMatchState;
 import com.google.cloud.modelarmor.v1.FilterResult;
+import com.google.cloud.modelarmor.v1.FloorSetting;
+import com.google.cloud.modelarmor.v1.FloorSettingName;
 import com.google.cloud.modelarmor.v1.LocationName;
 import com.google.cloud.modelarmor.v1.MaliciousUriFilterSettings;
 import com.google.cloud.modelarmor.v1.MaliciousUriFilterSettings.MaliciousUriFilterEnforcement;
@@ -46,6 +48,7 @@
 import com.google.cloud.modelarmor.v1.SdpFinding;
 import com.google.cloud.modelarmor.v1.Template;
 import com.google.cloud.modelarmor.v1.TemplateName;
+import com.google.cloud.modelarmor.v1.UpdateFloorSettingRequest;
 import com.google.privacy.dlp.v2.CreateDeidentifyTemplateRequest;
 import com.google.privacy.dlp.v2.CreateInspectTemplateRequest;
 import com.google.privacy.dlp.v2.DeidentifyConfig;
@@ -72,6 +75,7 @@
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -80,6 +84,10 @@
 public class SnippetsIT {
 
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
+  private static final String FOLDER_ID = System.getenv()
+      .getOrDefault("MA_FOLDER_ID", "global");
+  private static final String ORGANIZATION_ID = System.getenv()
+      .getOrDefault("MA_ORG_ID", "global");
   private static final String LOCATION_ID = System.getenv()
       .getOrDefault("GOOGLE_CLOUD_PROJECT_LOCATION", "us-central1");
   private static final String MA_ENDPOINT = String.format("modelarmor.%s.rep.googleapis.com:443",
@@ -99,7 +107,11 @@ public class SnippetsIT {
   private static String TEST_DEIDENTIFY_TEMPLATE_NAME;
   private ByteArrayOutputStream stdOut;
   private PrintStream originalOut;
+  private static String[] floorSettingNames;
   private static String[] templateToDelete;
+  private static String projectFloorSettingName;
+  private static String folderFloorSettingName;
+  private static String organizationFloorSettingName;
 
   // Check if the required environment variables are set.
   private static void requireEnvVar(String varName) {
@@ -111,6 +123,14 @@ private static void requireEnvVar(String varName) {
   @BeforeClass
   public static void beforeAll() throws IOException {
     requireEnvVar("GOOGLE_CLOUD_PROJECT");
+    requireEnvVar("MA_FOLDER_ID");
+    requireEnvVar("MA_ORG_ID");
+
+    projectFloorSettingName =
+        FloorSettingName.ofProjectLocationName(PROJECT_ID, "global").toString();
+    folderFloorSettingName = FloorSettingName.ofFolderLocationName(FOLDER_ID, "global").toString();
+    organizationFloorSettingName =
+        FloorSettingName.ofOrganizationLocationName(ORGANIZATION_ID, "global").toString();
 
     TEST_TEMPLATE_ID = randomId();
     TEST_RAI_TEMPLATE_ID = randomId();
@@ -147,6 +167,10 @@ private static String randomId() {
   @AfterClass
   public static void afterAll() throws IOException {
     requireEnvVar("GOOGLE_CLOUD_PROJECT");
+    requireEnvVar("MA_FOLDER_ID");
+    requireEnvVar("MA_ORG_ID");
+    
+    resetFloorSettings();
 
     // Delete templates after running tests.
     templateToDelete = new String[] {
@@ -380,6 +404,67 @@ private static void deleteTemplate(String templateId) throws IOException {
     }
   }
 
+  private static void resetFloorSettings() throws IOException {
+    floorSettingNames = new String[] {
+        projectFloorSettingName, folderFloorSettingName, organizationFloorSettingName
+    };
+    
+
+    try (ModelArmorClient client = ModelArmorClient.create()) {
+      for (String name : floorSettingNames) {
+        FloorSetting floorSetting = FloorSetting.newBuilder()
+            .setName(name)
+            .setFilterConfig(FilterConfig.newBuilder().build())
+            .setEnableFloorSettingEnforcement(false)
+            .build();
+
+        UpdateFloorSettingRequest request = UpdateFloorSettingRequest.newBuilder()
+            .setFloorSetting(floorSetting)
+            .build();
+
+        client.updateFloorSetting(request);
+      }
+    }
+  }
+
+  // Tests for Folder setting snippets.
+  @Test
+  public void testGetOrganizationFloorSetting() throws IOException {
+    GetOrganizationFloorSetting.getOrganizationFloorSetting(ORGANIZATION_ID);
+    assertThat(stdOut.toString()).contains("Fetched floor setting for organization:");
+  }
+  
+  @Test
+  public void testGetFolderFloorSetting() throws IOException {
+    GetFolderFloorSetting.getFolderFloorSetting(FOLDER_ID);
+    assertThat(stdOut.toString()).contains("Fetched floor setting for folder:");
+  }
+
+  @Test
+  public void testGetProjectFloorSetting() throws IOException {
+    GetProjectFloorSetting.getProjectFloorSetting(PROJECT_ID);
+    assertThat(stdOut.toString()).contains("Fetched floor setting for project:");
+  }
+
+  @Test
+  public void testUpdateOrganizationFloorSetting() throws IOException {
+    UpdateOrganizationsFloorSetting.updateOrganizationFloorSetting(ORGANIZATION_ID);
+    assertThat(stdOut.toString()).contains("Updated floor setting for organization:");
+  }
+
+  @Test
+  public void testUpdateFolderFloorSetting() throws IOException {
+    UpdateFolderFloorSetting.updateFolderFloorSetting(FOLDER_ID);
+    assertThat(stdOut.toString()).contains("Updated floor setting for folder:");
+  }
+
+
+  @Test
+  public void testUpdateProjectFloorSetting() throws IOException {
+    UpdateProjectFloorSetting.updateProjectFloorSetting(PROJECT_ID);
+    assertThat(stdOut.toString()).contains("Updated floor setting for project:");
+  }
+
   // Tests for Template CRUD snippets.
   @Test
   public void testUpdateModelArmorTemplate() throws IOException {