Skip to content

Commit d49bf96

Browse files
glasntglasnt
authored
Merge branch 'main' into renovate/kubernetes-go
2 parents 30bb385 + 4ff974f commit d49bf96

File tree

7 files changed

+394
-52
lines changed

7 files changed

+394
-52
lines changed

‎modelarmor/go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ require (
1010
github.com/googleapis/gax-go/v2 v2.14.1
1111
google.golang.org/api v0.228.0
1212
google.golang.org/grpc v1.71.0
13+
google.golang.org/protobuf v1.36.6
1314
)
1415

1516
require (
@@ -53,5 +54,4 @@ require (
5354
google.golang.org/genproto v0.0.0-20250313205543-e70fdf4c4cb4 // indirect
5455
google.golang.org/genproto/googleapis/api v0.0.0-20250313205543-e70fdf4c4cb4 // indirect
5556
google.golang.org/genproto/googleapis/rpc v0.0.0-20250313205543-e70fdf4c4cb4 // indirect
56-
google.golang.org/protobuf v1.36.6 // indirect
5757
)

‎modelarmor/modelarmor_test.go‎

Lines changed: 102 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ func testClient(t *testing.T) (*modelarmor.Client, context.Context) {
6969
return client, ctx
7070
}
7171

72-
// testModelArmorTemplate creates a new ModelArmor template for use in tests.
73-
// It returns the created template or an error.
72+
// testModelArmorTemplate creates a new Model Armor template with default
73+
// filter settings for use in integration tests. Returns the created template.
7474
func testModelArmorTemplate(t *testing.T, templateID string) (*modelarmorpb.Template, error) {
7575
t.Helper()
7676
tc := testutil.SystemTest(t)
@@ -103,29 +103,19 @@ func testModelArmorTemplate(t *testing.T, templateID string) (*modelarmorpb.Temp
103103
return response, nil
104104
}
105105

106-
// testModelArmorEmptyTemplate creates a new ModelArmor template for use in tests.
107-
// It returns the empty template or an error.
108-
func testModelArmorEmptyTemplate(t *testing.T, templateID string) (*modelarmorpb.Template, error) {
106+
// testCleanupTemplate deletes the specified Model Armor template if it exists,
107+
// ignoring the error if the template is already deleted.
108+
func testCleanupTemplate(t *testing.T, templateName string) {
109109
t.Helper()
110-
tc := testutil.SystemTest(t)
111-
locationID := testLocation(t)
112-
client, ctx := testClient(t)
113-
114-
template := &modelarmorpb.Template{
115-
FilterConfig: &modelarmorpb.FilterConfig{}}
116110

117-
req := &modelarmorpb.CreateTemplateRequest{
118-
Parent: fmt.Sprintf("projects/%s/locations/%s", tc.ProjectID, locationID),
119-
TemplateId: templateID,
120-
Template: template,
111+
client, ctx := testClient(t)
112+
err := client.DeleteTemplate(ctx, &modelarmorpb.DeleteTemplateRequest{Name: templateName})
113+
if err == nil {
114+
return
121115
}
122-
123-
response, err := client.CreateTemplate(ctx, req)
124-
if err != nil {
125-
return nil, fmt.Errorf("failed to create template: %w", err)
116+
if terr, ok := grpcstatus.FromError(err); !ok || terr.Code() != grpccodes.NotFound {
117+
t.Fatalf("testCleanupTemplate: failed to delete template %v", err)
126118
}
127-
128-
return response, nil
129119
}
130120

131121
// testAllFilterTemplate creates a new ModelArmor template with all filters enabled.
@@ -203,6 +193,31 @@ func testAllFilterTemplate(t *testing.T, templateID string) (*modelarmorpb.Templ
203193
return response, filterConfig, nil
204194
}
205195

196+
// testModelArmorEmptyTemplate creates a new ModelArmor template for use in tests.
197+
// It returns the empty template or an error.
198+
func testModelArmorEmptyTemplate(t *testing.T, templateID string) (*modelarmorpb.Template, error) {
199+
t.Helper()
200+
tc := testutil.SystemTest(t)
201+
locationID := testLocation(t)
202+
client, ctx := testClient(t)
203+
204+
template := &modelarmorpb.Template{
205+
FilterConfig: &modelarmorpb.FilterConfig{}}
206+
207+
req := &modelarmorpb.CreateTemplateRequest{
208+
Parent: fmt.Sprintf("projects/%s/locations/%s", tc.ProjectID, locationID),
209+
TemplateId: templateID,
210+
Template: template,
211+
}
212+
213+
response, err := client.CreateTemplate(ctx, req)
214+
if err != nil {
215+
return nil, fmt.Errorf("failed to create template: %w", err)
216+
}
217+
218+
return response, nil
219+
}
220+
206221
// testSDPTemplate creates DLP inspect and deidentify templates for use in tests.
207222
func testSDPTemplate(t *testing.T, projectID string, locationID string) (string, string) {
208223
inspectTemplateID := fmt.Sprintf("model-armor-inspect-template-%s", uuid.New().String())
@@ -383,21 +398,6 @@ func testModelArmorAdvancedSDPTemplate(t *testing.T, templateID string) (*modela
383398

384399
}
385400

386-
// testCleanupTemplate deletes the specified Model Armor template if it exists,
387-
// ignoring the error if the template is already deleted.
388-
func testCleanupTemplate(t *testing.T, templateName string) {
389-
t.Helper()
390-
391-
client, ctx := testClient(t)
392-
err := client.DeleteTemplate(ctx, &modelarmorpb.DeleteTemplateRequest{Name: templateName})
393-
if err == nil {
394-
return
395-
}
396-
if terr, ok := grpcstatus.FromError(err); !ok || terr.Code() != grpccodes.NotFound {
397-
t.Fatalf("testCleanupTemplate: failed to delete template %v", err)
398-
}
399-
}
400-
401401
// TestCreateModelArmorTemplateWithAdvancedSDP tests creating a
402402
// Model Armor template with advanced SDP using DLP templates.
403403
func TestCreateModelArmorTemplateWithAdvancedSDP(t *testing.T) {
@@ -1161,3 +1161,69 @@ func TestListModelArmorTemplatesWithFilter(t *testing.T) {
11611161
t.Errorf("listModelArmorTemplatesWithFilter: expected %q to contain %q", got, want)
11621162
}
11631163
}
1164+
1165+
// TestUpdateTemplate verifies that the updateModelArmorTemplate function
1166+
// successfully updates the filter configuration of an existing template.
1167+
func TestUpdateTemplate(t *testing.T) {
1168+
tc := testutil.SystemTest(t)
1169+
locationID := testLocation(t)
1170+
templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
1171+
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)
1172+
var buf bytes.Buffer
1173+
if _, err := testModelArmorTemplate(t, templateID); err != nil {
1174+
t.Fatal(err)
1175+
}
1176+
defer testCleanupTemplate(t, templateName)
1177+
1178+
if err := updateModelArmorTemplate(&buf, tc.ProjectID, locationID, templateID); err != nil {
1179+
t.Fatal(err)
1180+
}
1181+
1182+
if got, want := buf.String(), "Updated Filter Config: "; !strings.Contains(got, want) {
1183+
t.Errorf("updateModelArmorTemplate: expected %q to contain %q", got, want)
1184+
}
1185+
}
1186+
1187+
// TestUpdateTemplate verifies that the updateModelArmorTemplate function
1188+
// successfully updates the filter configuration of an existing template.
1189+
func TestUpdateTemplateLabels(t *testing.T) {
1190+
tc := testutil.SystemTest(t)
1191+
locationID := testLocation(t)
1192+
templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
1193+
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)
1194+
var buf bytes.Buffer
1195+
if _, err := testModelArmorTemplate(t, templateID); err != nil {
1196+
t.Fatal(err)
1197+
}
1198+
defer testCleanupTemplate(t, templateName)
1199+
1200+
if err := updateModelArmorTemplateLabels(&buf, tc.ProjectID, locationID, templateID, map[string]string{"testkey": "testvalue"}); err != nil {
1201+
t.Fatal(err)
1202+
}
1203+
1204+
if got, want := buf.String(), "Updated Model Armor Template Labels: "; !strings.Contains(got, want) {
1205+
t.Errorf("updateModelArmorTemplateLabels: expected %q to contain %q", got, want)
1206+
}
1207+
}
1208+
1209+
// TestUpdateTemplateMetadata verifies that the updateModelArmorTemplateMetadata function
1210+
// successfully updates the filter configuration of an existing template.
1211+
func TestUpdateTemplateMetadata(t *testing.T) {
1212+
tc := testutil.SystemTest(t)
1213+
locationID := testLocation(t)
1214+
templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
1215+
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)
1216+
var buf bytes.Buffer
1217+
if _, err := testModelArmorTemplate(t, templateID); err != nil {
1218+
t.Fatal(err)
1219+
}
1220+
defer testCleanupTemplate(t, templateName)
1221+
1222+
if err := updateModelArmorTemplateMetadata(&buf, tc.ProjectID, locationID, templateID); err != nil {
1223+
t.Fatal(err)
1224+
}
1225+
1226+
if got, want := buf.String(), "Updated Model Armor Template Metadata: "; !strings.Contains(got, want) {
1227+
t.Errorf("updateModelArmorTemplateMetadata: expected %q to contain %q", got, want)
1228+
}
1229+
}

‎modelarmor/update_template.go‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
// Sample code for updating the model armor template.
16+
17+
package modelarmor
18+
19+
// [START modelarmor_update_template]
20+
21+
import (
22+
"context"
23+
"fmt"
24+
"io"
25+
26+
modelarmor "cloud.google.com/go/modelarmor/apiv1"
27+
modelarmorpb "cloud.google.com/go/modelarmor/apiv1/modelarmorpb"
28+
"google.golang.org/api/option"
29+
)
30+
31+
// updateModelArmorTemplate updates a Model Armor template.
32+
//
33+
// updateModelArmorTemplate method updates a Model Armor template.
34+
//
35+
// w io.Writer: The writer to use for logging.
36+
// projectID string: The ID of the project.
37+
// locationID string: The ID of the location.
38+
// templateID string: The ID of the template.
39+
func updateModelArmorTemplate(w io.Writer, projectID, locationID, templateID string) error {
40+
ctx := context.Background()
41+
42+
// Create options for Model Armor client.
43+
opts := option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID))
44+
// Create the Model Armor client.
45+
client, err := modelarmor.NewClient(ctx, opts)
46+
if err != nil {
47+
return fmt.Errorf("failed to create client for project %s, location %s: %w", projectID, locationID, err)
48+
}
49+
defer client.Close()
50+
51+
// Build the Model Armor template with your preferred filters.
52+
// For more details on filters, please refer to the following doc:
53+
// [https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters](https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters)
54+
updatedTemplate := &modelarmorpb.Template{
55+
Name: fmt.Sprintf("projects/%s/locations/%s/templates/%s", projectID, locationID, templateID),
56+
FilterConfig: &modelarmorpb.FilterConfig{
57+
PiAndJailbreakFilterSettings: &modelarmorpb.PiAndJailbreakFilterSettings{
58+
FilterEnforcement: modelarmorpb.PiAndJailbreakFilterSettings_ENABLED,
59+
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_LOW_AND_ABOVE,
60+
},
61+
MaliciousUriFilterSettings: &modelarmorpb.MaliciousUriFilterSettings{
62+
FilterEnforcement: modelarmorpb.MaliciousUriFilterSettings_ENABLED,
63+
},
64+
},
65+
}
66+
67+
// Initialize request argument(s).
68+
req := &modelarmorpb.UpdateTemplateRequest{
69+
Template: updatedTemplate,
70+
}
71+
72+
// Update the template.
73+
response, err := client.UpdateTemplate(ctx, req)
74+
if err != nil {
75+
return fmt.Errorf("failed to update template: %w", err)
76+
}
77+
78+
// Print the updated filters in the template.
79+
fmt.Fprintf(w, "Updated Filter Config: %+v\n", response.FilterConfig)
80+
81+
return nil
82+
}
83+
84+
// [END modelarmor_update_template]

‎modelarmor/update_template_labels.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
// Sample code for updating the labels of the given model armor template.
16+
17+
package modelarmor
18+
19+
// [START modelarmor_update_template_with_labels]
20+
21+
import (
22+
"context"
23+
"fmt"
24+
"io"
25+
26+
modelarmor "cloud.google.com/go/modelarmor/apiv1"
27+
modelarmorpb "cloud.google.com/go/modelarmor/apiv1/modelarmorpb"
28+
"google.golang.org/api/option"
29+
"google.golang.org/protobuf/types/known/fieldmaskpb"
30+
)
31+
32+
// updateModelArmorTemplateLabels method updates
33+
// the labels of a Model Armor template.
34+
//
35+
// w io.Writer: The writer to use for logging.
36+
// projectID string: The ID of the project.
37+
// locationID string: The ID of the location.
38+
// templateID string: The ID of the template.
39+
// labels map[string]string: The updated labels.
40+
func updateModelArmorTemplateLabels(w io.Writer, projectID, locationID, templateID string, labels map[string]string) error {
41+
ctx := context.Background()
42+
43+
// Create options for Model Armor client.
44+
opts := option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID))
45+
// Create the Model Armor client.
46+
client, err := modelarmor.NewClient(ctx, opts)
47+
if err != nil {
48+
return fmt.Errorf("failed to create client for project %s, location %s: %w", projectID, locationID, err)
49+
}
50+
defer client.Close()
51+
52+
// Build the Model Armor template with your preferred filters.
53+
// For more details on filters, please refer to the following doc:
54+
// [https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters](https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters)
55+
template := &modelarmorpb.Template{
56+
Name: fmt.Sprintf("projects/%s/locations/%s/templates/%s", projectID, locationID, templateID),
57+
Labels: labels,
58+
}
59+
60+
// Prepare the request to update the template.
61+
updateMask := &fieldmaskpb.FieldMask{
62+
Paths: []string{"labels"},
63+
}
64+
65+
req := &modelarmorpb.UpdateTemplateRequest{
66+
Template: template,
67+
UpdateMask: updateMask,
68+
}
69+
70+
// Update the template.
71+
response, err := client.UpdateTemplate(ctx, req)
72+
if err != nil {
73+
return fmt.Errorf("failed to update template: %w", err)
74+
}
75+
76+
fmt.Fprintf(w, "Updated Model Armor Template Labels: %s\n", response.Name)
77+
78+
return nil
79+
}
80+
81+
// [END modelarmor_update_template_with_labels]

0 commit comments

Comments
 (0)