diff --git a/modelarmor/src/test/java/modelarmor/SnippetsIT.java b/modelarmor/src/test/java/modelarmor/SnippetsIT.java index ede9b6fa591..2b30d9a623f 100644 --- a/modelarmor/src/test/java/modelarmor/SnippetsIT.java +++ b/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 {