diff --git a/modelarmor/pom.xml b/modelarmor/pom.xml index 57d56f73f0f..ae557ad832b 100644 --- a/modelarmor/pom.xml +++ b/modelarmor/pom.xml @@ -43,7 +43,7 @@ com.google.cloud libraries-bom - 26.59.0 + 26.60.0 pom import 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..0b6527857c5 --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateFolderFloorSetting.java @@ -0,0 +1,93 @@ +/* + * 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.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.LOW_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .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..5cb1d34b652 --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateOrganizationsFloorSetting.java @@ -0,0 +1,96 @@ +/* + * 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.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.LOW_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .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..ebe1eebda0a --- /dev/null +++ b/modelarmor/src/main/java/modelarmor/UpdateProjectFloorSetting.java @@ -0,0 +1,93 @@ +/* + * 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.HARASSMENT) + .setConfidenceLevel(DetectionConfidenceLevel.LOW_AND_ABOVE) + .build(), + RaiFilter.newBuilder() + .setFilterType(RaiFilterType.SEXUALLY_EXPLICIT) + .setConfidenceLevel(DetectionConfidenceLevel.HIGH) + .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..ede9b6fa591 100644 --- a/modelarmor/src/test/java/modelarmor/SnippetsIT.java +++ b/modelarmor/src/test/java/modelarmor/SnippetsIT.java @@ -84,6 +84,7 @@ public class SnippetsIT { .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 TEST_TEMPLATE_ID; private static String TEST_RAI_TEMPLATE_ID; private static String TEST_CSAM_TEMPLATE_ID; @@ -97,15 +98,14 @@ 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[] templateToDelete; // Check if the required environment variables are set. - private static String requireEnvVar(String varName) { - String value = System.getenv(varName); + private static void requireEnvVar(String varName) { assertNotNull( "Environment variable " + varName + " is required to run these tests.", System.getenv(varName)); - return value; } @BeforeClass @@ -139,6 +139,11 @@ 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"); @@ -162,6 +167,7 @@ public static void afterAll() throws IOException { @Before public void beforeEach() { + originalOut = System.out; stdOut = new ByteArrayOutputStream(); System.setOut(new PrintStream(stdOut)); } @@ -174,17 +180,12 @@ 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(); } - // Create Model Armor templates required for tests. - private static Template createMaliciousUriTemplate() throws IOException { + // Helper functions to manage templates. + private static void createMaliciousUriTemplate() throws IOException { // Create a malicious URI filter template. MaliciousUriFilterSettings maliciousUriFilterSettings = MaliciousUriFilterSettings.newBuilder() .setFilterEnforcement(MaliciousUriFilterEnforcement.ENABLED) @@ -199,10 +200,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 +220,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,128 +240,6 @@ private static Template createBasicSdpTemplate() throws IOException { .build(); createTemplate(template, TEST_BASIC_SDP_TEMPLATE_ID); - return template; - } - - @Test - public void testUpdateModelArmorTemplate() throws IOException { - CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - - // Update the existing template. - Template updatedTemplate = UpdateTemplate.updateTemplate(PROJECT_ID, LOCATION_ID, - TEST_TEMPLATE_ID); - - assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); - } - - @Test - public void testUpdateModelArmorTemplateWithLabels() throws IOException { - CreateTemplateWithLabels.createTemplateWithLabels(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - - // Update the existing template. - Template updatedTemplate = UpdateTemplateWithLabels.updateTemplateWithLabels(PROJECT_ID, - LOCATION_ID, TEST_TEMPLATE_ID); - - assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); - } - - @Test - public void testUpdateModelArmorTemplateWithMetadata() throws IOException { - CreateTemplateWithMetadata.createTemplateWithMetadata(PROJECT_ID, LOCATION_ID, - TEST_TEMPLATE_ID); - - // Update the existing template. - Template updatedTemplate = UpdateTemplateWithMetadata.updateTemplateWithMetadata(PROJECT_ID, - 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()); - } - - @Test - public void testGetModelArmorTemplate() throws IOException { - CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - Template retrievedTemplate = GetTemplate.getTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - - assertEquals(retrievedTemplate.getName(), TEST_TEMPLATE_NAME); - } - - @Test - public void testListModelArmorTemplates() throws IOException { - CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - - ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID); - - boolean templatePresentInList = false; - for (Template template : ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID).iterateAll()) { - if (TEST_TEMPLATE_NAME.equals(template.getName())) { - templatePresentInList = true; - } - } - assertTrue(templatePresentInList); - } - - @Test - public void testListTemplatesWithFilter() throws IOException { - CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - String filter = "name=\"projects/" + PROJECT_ID + "/locations/" + LOCATION_ID + "/" - + TEST_TEMPLATE_ID + "\""; - - ListTemplatesWithFilter.listTemplatesWithFilter(PROJECT_ID, LOCATION_ID, filter); - - boolean templatePresentInList = false; - for (Template template : ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID).iterateAll()) { - if (TEST_TEMPLATE_NAME.equals(template.getName())) { - templatePresentInList = true; - } - } - assertTrue(templatePresentInList); - } - - public void testCreateModelArmorTemplate() throws IOException { - Template createdTemplate = CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, - TEST_TEMPLATE_ID); - - assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); - } - - @Test - public void testCreateModelArmorTemplateWithBasicSDP() throws IOException { - Template createdTemplate = CreateTemplateWithBasicSdp.createTemplateWithBasicSdp(PROJECT_ID, - LOCATION_ID, TEST_TEMPLATE_ID); - - assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); - assertEquals(SdpBasicConfigEnforcement.ENABLED, - createdTemplate.getFilterConfig().getSdpSettings().getBasicConfig().getFilterEnforcement()); - } - - @Test - public void testCreateModelArmorTemplateWithLabels() throws IOException { - Template createdTemplate = CreateTemplateWithLabels.createTemplateWithLabels(PROJECT_ID, - LOCATION_ID, TEST_TEMPLATE_ID); - - assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); - } - - @Test - public void testCreateModelArmorTemplateWithMetadata() throws IOException { - Template createdTemplate = CreateTemplateWithMetadata.createTemplateWithMetadata(PROJECT_ID, - LOCATION_ID, TEST_TEMPLATE_ID); - - assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); - assertEquals(true, createdTemplate.getTemplateMetadata().getIgnorePartialInvocationFailures()); - assertEquals(true, createdTemplate.getTemplateMetadata().getLogSanitizeOperations()); - assertEquals(500, createdTemplate.getTemplateMetadata().getCustomPromptSafetyErrorCode()); - } - - @Test - public void testDeleteModelArmorTemplate() throws IOException { - CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - DeleteTemplate.deleteTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); - - assertThat(stdOut.toString()).contains("Deleted template:"); } private static void deleteModelArmorTemplate(String templateId) throws IOException { @@ -386,10 +263,10 @@ private static InspectTemplate createInspectTemplate(String templateId) throws I try (DlpServiceClient dlpServiceClient = DlpServiceClient.create()) { // Info Types: // https://cloud.google.com/sensitive-data-protection/docs/infotypes-reference - List infoTypes = Stream - .of("PHONE_NUMBER", "EMAIL_ADDRESS", "US_INDIVIDUAL_TAXPAYER_IDENTIFICATION_NUMBER") - .map(it -> InfoType.newBuilder().setName(it).build()) - .collect(Collectors.toList()); + List infoTypes = + Stream.of("PHONE_NUMBER", "EMAIL_ADDRESS", "US_INDIVIDUAL_TAXPAYER_IDENTIFICATION_NUMBER") + .map(it -> InfoType.newBuilder().setName(it).build()) + .collect(Collectors.toList()); InspectConfig inspectConfig = InspectConfig.newBuilder() .addAllInfoTypes(infoTypes) @@ -399,13 +276,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); } @@ -503,6 +380,144 @@ private static void deleteTemplate(String templateId) throws IOException { } } + // Tests for Template CRUD snippets. + @Test + public void testUpdateModelArmorTemplate() throws IOException { + CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + + // Update the existing template. + Template updatedTemplate = UpdateTemplate.updateTemplate(PROJECT_ID, LOCATION_ID, + TEST_TEMPLATE_ID); + + assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); + } + + @Test + public void testUpdateModelArmorTemplateWithLabels() throws IOException { + CreateTemplateWithLabels.createTemplateWithLabels(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + + // Update the existing template. + Template updatedTemplate = UpdateTemplateWithLabels.updateTemplateWithLabels(PROJECT_ID, + LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); + } + + @Test + public void testUpdateModelArmorTemplateWithMetadata() throws IOException { + CreateTemplateWithMetadata.createTemplateWithMetadata(PROJECT_ID, LOCATION_ID, + TEST_TEMPLATE_ID); + + // Update the existing template. + Template updatedTemplate = UpdateTemplateWithMetadata.updateTemplateWithMetadata(PROJECT_ID, + LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(updatedTemplate.getName(), TEST_TEMPLATE_NAME); + assertEquals(true, updatedTemplate.getTemplateMetadata().getLogTemplateOperations()); + assertEquals(true, updatedTemplate.getTemplateMetadata().getLogSanitizeOperations()); + } + + @Test + public void testGetModelArmorTemplate() throws IOException { + CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + Template retrievedTemplate = GetTemplate.getTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(retrievedTemplate.getName(), TEST_TEMPLATE_NAME); + } + + @Test + public void testListModelArmorTemplates() throws IOException { + CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + + ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID); + + boolean templatePresentInList = false; + for (Template template : ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID).iterateAll()) { + if (TEST_TEMPLATE_NAME.equals(template.getName())) { + templatePresentInList = true; + } + } + assertTrue(templatePresentInList); + } + + @Test + public void testListTemplatesWithFilter() throws IOException { + CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + String filter = "name=\"projects/" + PROJECT_ID + "/locations/" + LOCATION_ID + "/" + + TEST_TEMPLATE_ID + "\""; + + ListTemplatesWithFilter.listTemplatesWithFilter(PROJECT_ID, LOCATION_ID, filter); + + boolean templatePresentInList = false; + for (Template template : ListTemplates.listTemplates(PROJECT_ID, LOCATION_ID).iterateAll()) { + if (TEST_TEMPLATE_NAME.equals(template.getName())) { + templatePresentInList = true; + } + } + assertTrue(templatePresentInList); + } + + @Test + public void testCreateModelArmorTemplate() throws IOException { + Template createdTemplate = CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, + TEST_TEMPLATE_ID); + + assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); + } + + @Test + public void testCreateModelArmorTemplateWithBasicSDP() throws IOException { + Template createdTemplate = CreateTemplateWithBasicSdp.createTemplateWithBasicSdp(PROJECT_ID, + LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); + assertEquals(SdpBasicConfigEnforcement.ENABLED, + createdTemplate.getFilterConfig().getSdpSettings().getBasicConfig().getFilterEnforcement()); + } + + @Test + public void testCreateModelArmorTemplateWithAdvancedSDP() throws IOException { + + Template createdTemplate = CreateTemplateWithAdvancedSdp.createTemplateWithAdvancedSdp( + PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID, + TEST_INSPECT_TEMPLATE_ID, TEST_DEIDENTIFY_TEMPLATE_ID); + + assertEquals(TEST_TEMPLATE_NAME, createdTemplate.getName()); + + SdpAdvancedConfig advancedSdpConfig = createdTemplate.getFilterConfig().getSdpSettings() + .getAdvancedConfig(); + + assertEquals(TEST_INSPECT_TEMPLATE_NAME, advancedSdpConfig.getInspectTemplate()); + assertEquals(TEST_DEIDENTIFY_TEMPLATE_NAME, advancedSdpConfig.getDeidentifyTemplate()); + } + + @Test + public void testCreateModelArmorTemplateWithLabels() throws IOException { + Template createdTemplate = CreateTemplateWithLabels.createTemplateWithLabels(PROJECT_ID, + LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); + } + + @Test + public void testCreateModelArmorTemplateWithMetadata() throws IOException { + Template createdTemplate = CreateTemplateWithMetadata.createTemplateWithMetadata(PROJECT_ID, + LOCATION_ID, TEST_TEMPLATE_ID); + + assertEquals(createdTemplate.getName(), TEST_TEMPLATE_NAME); + assertEquals(true, createdTemplate.getTemplateMetadata().getLogTemplateOperations()); + assertEquals(true, createdTemplate.getTemplateMetadata().getLogSanitizeOperations()); + } + + @Test + public void testDeleteModelArmorTemplate() throws IOException { + CreateTemplate.createTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + DeleteTemplate.deleteTemplate(PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID); + + assertThat(stdOut.toString()).contains("Deleted template:"); + } + + // Tests for user prompt sanitization snippets. @Test public void testSanitizeUserPromptWithRaiTemplate() throws IOException { String userPrompt = "How to make cheesecake without oven at home?"; @@ -675,6 +690,7 @@ public void testSanitizeUserPromptWithAdvancedSdpTemplate() throws IOException { } } + // Tests for model response sanitization snippets. @Test public void testSanitizeModelResponseWithRaiTemplate() throws IOException { String modelResponse = "To make cheesecake without oven, you'll need to follow these steps..."; @@ -704,6 +720,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"; @@ -836,20 +853,4 @@ public void testScreenPdfFile() throws IOException { assertEquals(FilterMatchState.NO_MATCH_FOUND, response.getSanitizationResult().getFilterMatchState()); } - - @Test - public void testCreateModelArmorTemplateWithAdvancedSDP() throws IOException { - - Template createdTemplate = CreateTemplateWithAdvancedSdp.createTemplateWithAdvancedSdp( - PROJECT_ID, LOCATION_ID, TEST_TEMPLATE_ID, - TEST_INSPECT_TEMPLATE_ID, TEST_DEIDENTIFY_TEMPLATE_ID); - - assertEquals(TEST_TEMPLATE_NAME, createdTemplate.getName()); - - SdpAdvancedConfig advancedSdpConfig = createdTemplate.getFilterConfig().getSdpSettings() - .getAdvancedConfig(); - - assertEquals(TEST_INSPECT_TEMPLATE_NAME, advancedSdpConfig.getInspectTemplate()); - assertEquals(TEST_DEIDENTIFY_TEMPLATE_NAME, advancedSdpConfig.getDeidentifyTemplate()); - } }