diff --git a/modelarmor/src/main/java/modelarmor/CreateTemplateWithMetadata.java b/modelarmor/src/main/java/modelarmor/CreateTemplateWithMetadata.java index b6e320b4078..c70de6c1f1e 100644 --- a/modelarmor/src/main/java/modelarmor/CreateTemplateWithMetadata.java +++ b/modelarmor/src/main/java/modelarmor/CreateTemplateWithMetadata.java @@ -95,9 +95,8 @@ public static Template createTemplateWithMetadata( // For more details about metadata, refer to the following documentation: // https://cloud.google.com/security-command-center/docs/reference/model-armor/rest/v1/projects.locations.templates#templatemetadata TemplateMetadata templateMetadata = TemplateMetadata.newBuilder() - .setIgnorePartialInvocationFailures(true) + .setLogTemplateOperations(true) .setLogSanitizeOperations(true) - .setCustomPromptSafetyErrorCode(500) .build(); Template template = Template.newBuilder() diff --git a/modelarmor/src/main/java/modelarmor/GetFolderFloorSetting.java b/modelarmor/src/main/java/modelarmor/GetFolderFloorSetting.java new file mode 100644 index 00000000000..ba44bf3ee6a --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/GetFolderFloorSetting.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_get_folder_floor_settings] + +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.GetFloorSettingRequest; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import java.io.IOException; + +public class GetFolderFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String folderId = "your-folder-id"; + + getFolderFloorSetting(folderId); + } + + public static FloorSetting getFolderFloorSetting(String folderId) throws IOException { + + // 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(); + + GetFloorSettingRequest request = GetFloorSettingRequest.newBuilder().setName(name).build(); + + FloorSetting floorSetting = client.getFloorSetting(request); + System.out.println("Fetched floor setting for folder: " + folderId); + + return floorSetting; + } + } +} +// [END modelarmor_get_folder_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/GetOrganizationFloorSetting.java b/modelarmor/src/main/java/modelarmor/GetOrganizationFloorSetting.java new file mode 100644 index 00000000000..d010e89f580 --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/GetOrganizationFloorSetting.java @@ -0,0 +1,53 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_get_organization_floor_settings] + +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.GetFloorSettingRequest; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import java.io.IOException; + +public class GetOrganizationFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String organizationId = "your-organization-id"; + + getOrganizationFloorSetting(organizationId); + } + + public static FloorSetting getOrganizationFloorSetting(String organizationId) throws IOException { + + // 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.ofOrganizationLocationName(organizationId, "global") + .toString(); + + GetFloorSettingRequest request = GetFloorSettingRequest.newBuilder().setName(name).build(); + + FloorSetting floorSetting = client.getFloorSetting(request); + System.out.println("Fetched floor setting for organization: " + organizationId); + + return floorSetting; + } + } +} +// [END modelarmor_get_organization_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/GetProjectFloorSetting.java b/modelarmor/src/main/java/modelarmor/GetProjectFloorSetting.java new file mode 100644 index 00000000000..84bf669deea --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/GetProjectFloorSetting.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_get_project_floor_settings] + +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.GetFloorSettingRequest; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import java.io.IOException; + +public class GetProjectFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String projectId = "your-project-id"; + + getProjectFloorSetting(projectId); + } + + public static FloorSetting getProjectFloorSetting(String projectId) throws IOException { + + // 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(projectId, "global").toString(); + + GetFloorSettingRequest request = GetFloorSettingRequest.newBuilder().setName(name).build(); + + FloorSetting floorSetting = client.getFloorSetting(request); + System.out.println("Fetched floor setting for project: " + projectId); + + return floorSetting; + } + } +} +// [END modelarmor_get_project_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/UpdateFolderFloorSetting.java b/modelarmor/src/main/java/modelarmor/UpdateFolderFloorSetting.java new file mode 100644 index 00000000000..6e347d355eb --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateFolderFloorSetting.java @@ -0,0 +1,101 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_update_folder_floor_settings] + +import com.google.cloud.modelarmor.v1.DetectionConfidenceLevel; +import com.google.cloud.modelarmor.v1.FilterConfig; +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import com.google.cloud.modelarmor.v1.RaiFilterSettings; +import com.google.cloud.modelarmor.v1.RaiFilterSettings.RaiFilter; +import com.google.cloud.modelarmor.v1.RaiFilterType; +import com.google.cloud.modelarmor.v1.UpdateFloorSettingRequest; +import com.google.protobuf.FieldMask; +import java.io.IOException; +import java.util.List; + +public class UpdateFolderFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String folderId = "your-folder-id"; + + updateFolderFloorSetting(folderId); + } + + public static FloorSetting updateFolderFloorSetting(String folderId) + throws IOException { + + // 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.ofFolderLocationName(folderId, "global").toString(); + + // For more details on filters, please refer to the following doc: + // https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters + RaiFilterSettings raiFilterSettings = + RaiFilterSettings.newBuilder() + .addAllRaiFilters( + List.of( + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.DANGEROUS) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HATE_SPEECH) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build())) + .build(); + + FilterConfig modelArmorFilter = FilterConfig.newBuilder() + .setRaiSettings(raiFilterSettings) + .build(); + + // Create a field mask to specify which fields to update. + // Ref: https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask + FieldMask updateMask = FieldMask.newBuilder().addPaths("filter_config.rai_settings").build(); + + FloorSetting floorSetting = FloorSetting.newBuilder() + .setName(name) + .setFilterConfig(modelArmorFilter) + .setEnableFloorSettingEnforcement(true) + .build(); + + UpdateFloorSettingRequest request = UpdateFloorSettingRequest.newBuilder() + .setFloorSetting(floorSetting) + .setUpdateMask(updateMask) + .build(); + + FloorSetting updatedFloorSetting = client.updateFloorSetting(request); + System.out.println("Updated floor setting for folder: " + folderId); + + return updatedFloorSetting; + } + } +} +// [END modelarmor_update_folder_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/UpdateOrganizationsFloorSetting.java b/modelarmor/src/main/java/modelarmor/UpdateOrganizationsFloorSetting.java new file mode 100644 index 00000000000..9f7728b0ee4 --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateOrganizationsFloorSetting.java @@ -0,0 +1,104 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_update_organization_floor_settings] + +import com.google.cloud.modelarmor.v1.DetectionConfidenceLevel; +import com.google.cloud.modelarmor.v1.FilterConfig; +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import com.google.cloud.modelarmor.v1.RaiFilterSettings; +import com.google.cloud.modelarmor.v1.RaiFilterSettings.RaiFilter; +import com.google.cloud.modelarmor.v1.RaiFilterType; +import com.google.cloud.modelarmor.v1.UpdateFloorSettingRequest; +import com.google.protobuf.FieldMask; +import java.io.IOException; +import java.util.List; + +public class UpdateOrganizationsFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String organizationId = "your-organization-id"; + + updateOrganizationFloorSetting(organizationId); + } + + public static FloorSetting updateOrganizationFloorSetting(String organizationId) + throws IOException { + + // 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.ofOrganizationLocationName(organizationId, "global") + .toString(); + + // For more details on filters, please refer to the following doc: + // https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters + RaiFilterSettings raiFilterSettings = + RaiFilterSettings.newBuilder() + .addAllRaiFilters( + List.of( + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.DANGEROUS) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HATE_SPEECH) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build())) + .build(); + + FilterConfig modelArmorFilter = FilterConfig.newBuilder() + .setRaiSettings(raiFilterSettings) + .build(); + + // Create a field mask to specify which fields to update. + // Ref: https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask + FieldMask updateMask = FieldMask.newBuilder() + .addPaths("filter_config.rai_settings") + .build(); + + FloorSetting floorSetting = FloorSetting.newBuilder() + .setName(name) + .setFilterConfig(modelArmorFilter) + .setEnableFloorSettingEnforcement(true) + .build(); + + UpdateFloorSettingRequest request = UpdateFloorSettingRequest.newBuilder() + .setFloorSetting(floorSetting) + .setUpdateMask(updateMask) + .build(); + + FloorSetting updatedFloorSetting = client.updateFloorSetting(request); + System.out.println("Updated floor setting for organization: " + organizationId); + + return updatedFloorSetting; + } + } +} +// [END modelarmor_update_organization_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/UpdateProjectFloorSetting.java b/modelarmor/src/main/java/modelarmor/UpdateProjectFloorSetting.java new file mode 100644 index 00000000000..2e823036609 --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateProjectFloorSetting.java @@ -0,0 +1,103 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. +*/ + +package modelarmor; + +// [START modelarmor_update_project_floor_settings] + +import com.google.cloud.modelarmor.v1.DetectionConfidenceLevel; +import com.google.cloud.modelarmor.v1.FilterConfig; +import com.google.cloud.modelarmor.v1.FloorSetting; +import com.google.cloud.modelarmor.v1.FloorSettingName; +import com.google.cloud.modelarmor.v1.ModelArmorClient; +import com.google.cloud.modelarmor.v1.RaiFilterSettings; +import com.google.cloud.modelarmor.v1.RaiFilterSettings.RaiFilter; +import com.google.cloud.modelarmor.v1.RaiFilterType; +import com.google.cloud.modelarmor.v1.UpdateFloorSettingRequest; +import com.google.protobuf.FieldMask; +import java.io.IOException; +import java.util.List; + +public class UpdateProjectFloorSetting { + + public static void main(String[] args) throws IOException { + // TODO(developer): Replace these variables before running the sample. + String projectId = "your-project-id"; + + updateProjectFloorSetting(projectId); + } + + public static FloorSetting updateProjectFloorSetting(String projectId) + throws IOException { + + // 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(projectId, "global").toString(); + + // For more details on filters, please refer to the following doc: + // https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters + RaiFilterSettings raiFilterSettings = + RaiFilterSettings.newBuilder() + .addAllRaiFilters( + List.of( + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.DANGEROUS) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HATE_SPEECH) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.MEDIUM_AND_ABOVE) + .build())) + .build(); + + FilterConfig modelArmorFilter = FilterConfig.newBuilder() + .setRaiSettings(raiFilterSettings) + .build(); + + // Create a field mask to specify which fields to update. + // Ref: https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask + FieldMask updateMask = FieldMask.newBuilder() + .addPaths("filter_config.rai_settings") + .build(); + + FloorSetting floorSetting = FloorSetting.newBuilder() + .setName(name) + .setFilterConfig(modelArmorFilter) + .setEnableFloorSettingEnforcement(true) + .build(); + + UpdateFloorSettingRequest request = UpdateFloorSettingRequest.newBuilder() + .setFloorSetting(floorSetting) + .setUpdateMask(updateMask) + .build(); + + FloorSetting updatedFloorSetting = client.updateFloorSetting(request); + System.out.println("Updated floor setting for project: " + projectId); + + return updatedFloorSetting; + } + } +} +// [END modelarmor_update_project_floor_settings] diff --git a/modelarmor/src/main/java/modelarmor/UpdateTemplateWithMetadata.java b/modelarmor/src/main/java/modelarmor/UpdateTemplateWithMetadata.java index 8fc3563d3be..7fff2bc16b9 100644 --- a/modelarmor/src/main/java/modelarmor/UpdateTemplateWithMetadata.java +++ b/modelarmor/src/main/java/modelarmor/UpdateTemplateWithMetadata.java @@ -59,9 +59,8 @@ public static Template updateTemplateWithMetadata(String projectId, String locat // For more details about metadata, refer to the following documentation: // https://cloud.google.com/security-command-center/docs/reference/model-armor/rest/v1/projects.locations.templates#templatemetadata TemplateMetadata updatedMetadata = TemplateMetadata.newBuilder() - .setIgnorePartialInvocationFailures(false) - .setLogSanitizeOperations(false) - .setCustomPromptSafetyErrorCode(400) + .setLogTemplateOperations(true) + .setLogSanitizeOperations(true) .build(); // Update the template with new metadata. diff --git a/modelarmor/src/test/java/modelarmor/SnippetsIT.java b/modelarmor/src/test/java/modelarmor/SnippetsIT.java index 35f7d69e5ad..1be7b73b44d 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; @@ -80,10 +83,20 @@ public class SnippetsIT { private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT"); + private static final String FOLDER_ID = System.getenv("MA_FOLDER_ID"); + private static final String ORGANIZATION_ID = System.getenv("MA_ORG_ID"); 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", LOCATION_ID); + + private static String projectFloorSettingName = FloorSettingName + .ofProjectLocationName(PROJECT_ID, "global").toString(); + private static String folderFloorSettingName = FloorSettingName + .ofFolderLocationName(FOLDER_ID, "global").toString(); + private static String organizationFloorSettingName = FloorSettingName + .ofOrganizationLocationName(ORGANIZATION_ID, "global").toString(); + private static String TEST_TEMPLATE_ID; private static String TEST_RAI_TEMPLATE_ID; private static String TEST_CSAM_TEMPLATE_ID; @@ -97,6 +110,8 @@ public class SnippetsIT { private static String TEST_INSPECT_TEMPLATE_NAME; private static String TEST_DEIDENTIFY_TEMPLATE_NAME; private ByteArrayOutputStream stdOut; + private PrintStream originalOut; + private static String[] floorSettingNames; private static String[] templateToDelete; // Check if the required environment variables are set. @@ -111,6 +126,8 @@ private static String requireEnvVar(String varName) { @BeforeClass public static void beforeAll() throws IOException { requireEnvVar("GOOGLE_CLOUD_PROJECT"); + requireEnvVar("MA_FOLDER_ID"); + requireEnvVar("MA_ORG_ID"); TEST_TEMPLATE_ID = randomId(); TEST_RAI_TEMPLATE_ID = randomId(); @@ -139,9 +156,18 @@ public static void beforeAll() throws IOException { CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_CSAM_TEMPLATE_ID); } + private static String randomId() { + Random random = new Random(); + return "java-ma-" + random.nextLong(); + } + @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[] { @@ -162,6 +188,7 @@ public static void afterAll() throws IOException { @Before public void beforeEach() { + originalOut = System.out; stdOut = new ByteArrayOutputStream(); System.setOut(new PrintStream(stdOut)); } @@ -174,17 +201,70 @@ public void afterEach() throws IOException { // Ignore not found error - template already deleted. } + System.setOut(originalOut); stdOut = null; - System.setOut(null); } - private static String randomId() { - Random random = new Random(); - return "java-ma-" + random.nextLong(); + 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); + } + } + } + + @Test + public void testGetFolderFloorSetting() throws IOException { + GetFolderFloorSetting.getFolderFloorSetting(FOLDER_ID); + assertThat(stdOut.toString()).contains("Fetched floor setting for folder:"); + } + + @Test + public void testGetOrganizationFloorSetting() throws IOException { + GetOrganizationFloorSetting.getOrganizationFloorSetting(ORGANIZATION_ID); + assertThat(stdOut.toString()).contains("Fetched floor setting for organization:"); + } + + @Test + public void testGetProjectFloorSetting() throws IOException { + GetProjectFloorSetting.getProjectFloorSetting(PROJECT_ID); + assertThat(stdOut.toString()).contains("Fetched floor setting for project:"); + } + + @Test + public void testUpdateFolderFloorSetting() throws IOException { + UpdateFolderFloorSetting.updateFolderFloorSetting(FOLDER_ID); + assertThat(stdOut.toString()).contains("Updated floor setting for folder:"); + } + + @Test + public void testUpdateOrganizationFloorSetting() throws IOException { + UpdateOrganizationsFloorSetting.updateOrganizationFloorSetting(ORGANIZATION_ID); + assertThat(stdOut.toString()).contains("Updated floor setting for organization:"); + } + + @Test + public void testUpdateProjectFloorSetting() throws IOException { + UpdateProjectFloorSetting.updateProjectFloorSetting(PROJECT_ID); + assertThat(stdOut.toString()).contains("Updated floor setting for project:"); } // Create Model Armor templates required for tests. - private static Template createMaliciousUriTemplate() throws IOException { + private static void createMaliciousUriTemplate() throws IOException { // Create a malicious URI filter template. MaliciousUriFilterSettings maliciousUriFilterSettings = MaliciousUriFilterSettings.newBuilder() .setFilterEnforcement(MaliciousUriFilterEnforcement.ENABLED) @@ -199,10 +279,9 @@ private static Template createMaliciousUriTemplate() throws IOException { .build(); createTemplate(template, TEST_MALICIOUS_URI_TEMPLATE_ID); - return template; } - private static Template createPiAndJailBreakTemplate() throws IOException { + private static void createPiAndJailBreakTemplate() throws IOException { // Create a Pi and Jailbreak filter template. // Create a template with Prompt injection & Jailbreak settings. PiAndJailbreakFilterSettings piAndJailbreakFilterSettings = PiAndJailbreakFilterSettings @@ -220,10 +299,9 @@ private static Template createPiAndJailBreakTemplate() throws IOException { .build(); createTemplate(template, TEST_PI_JAILBREAK_TEMPLATE_ID); - return template; } - private static Template createBasicSdpTemplate() throws IOException { + private static void createBasicSdpTemplate() throws IOException { SdpBasicConfig basicSdpConfig = SdpBasicConfig.newBuilder() .setFilterEnforcement(SdpBasicConfigEnforcement.ENABLED) .build(); @@ -241,7 +319,6 @@ private static Template createBasicSdpTemplate() throws IOException { .build(); createTemplate(template, TEST_BASIC_SDP_TEMPLATE_ID); - return template; } @Test @@ -276,9 +353,8 @@ public void testUpdateModelArmorTemplateWithMetadata() throws IOException { LOCATION_ID, TEST_TEMPLATE_ID); assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); - assertEquals(false, updatedTemplate.getTemplateMetadata().getIgnorePartialInvocationFailures()); - assertEquals(false, updatedTemplate.getTemplateMetadata().getLogSanitizeOperations()); - assertEquals(400, updatedTemplate.getTemplateMetadata().getCustomPromptSafetyErrorCode()); + assertEquals(true, updatedTemplate.getTemplateMetadata().getLogTemplateOperations()); + assertEquals(true, updatedTemplate.getTemplateMetadata().getLogSanitizeOperations()); } @Test @@ -321,6 +397,7 @@ public void testListTemplatesWithFilter() throws IOException { assertTrue(templatePresentInList); } + @Test public void testCreateModelArmorTemplate() throws IOException { Template createdTemplate = CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); @@ -352,9 +429,8 @@ public void testCreateModelArmorTemplateWithMetadata() throws IOException { LOCATION_ID, TEST_TEMPLATE_ID); assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); - assertEquals(true, createdTemplate.getTemplateMetadata().getIgnorePartialInvocationFailures()); + assertEquals(true, createdTemplate.getTemplateMetadata().getLogTemplateOperations()); assertEquals(true, createdTemplate.getTemplateMetadata().getLogSanitizeOperations()); - assertEquals(500, createdTemplate.getTemplateMetadata().getCustomPromptSafetyErrorCode()); } @Test @@ -399,13 +475,13 @@ private static InspectTemplate createInspectTemplate(String templateId) throws I .setInspectConfig(inspectConfig) .build(); - CreateInspectTemplateRequest createInspectTemplateRequest = - CreateInspectTemplateRequest.newBuilder() - .setParent( - com.google.privacy.dlp.v2.LocationName.of(PROJECT_ID, LOCATION_ID).toString()) - .setTemplateId(templateId) - .setInspectTemplate(inspectTemplate) - .build(); + CreateInspectTemplateRequest createInspectTemplateRequest = CreateInspectTemplateRequest + .newBuilder() + .setParent( + com.google.privacy.dlp.v2.LocationName.of(PROJECT_ID, LOCATION_ID).toString()) + .setTemplateId(templateId) + .setInspectTemplate(inspectTemplate) + .build(); return dlpServiceClient.createInspectTemplate(createInspectTemplateRequest); } @@ -704,6 +780,7 @@ public void testSanitizeModelResponseWithRaiTemplate() throws IOException { } } + @Test public void testSanitizeModelResponseWithMaliciousUrlTemplate() throws IOException { String modelResponse = "You can use this to make a cake: https://testsafebrowsing.appspot.com/s/malware.html";