@@ -21,7 +21,10 @@ import (
2121 "os"
2222 "strings"
2323 "testing"
24+ "time"
2425
26+ dlp "cloud.google.com/go/dlp/apiv2"
27+ "cloud.google.com/go/dlp/apiv2/dlppb"
2528 modelarmor "cloud.google.com/go/modelarmor/apiv1"
2629 modelarmorpb "cloud.google.com/go/modelarmor/apiv1/modelarmorpb"
2730
@@ -78,6 +81,106 @@ func testCleanupTemplate(t *testing.T, templateName string) {
7881 }
7982}
8083
84+ // testSDPTemplate creates DLP inspect and deidentify templates for use in tests.
85+ func testSDPTemplate (t * testing.T , projectID string , locationID string ) (string , string ) {
86+ inspectTemplateID := fmt .Sprintf ("model-armor-inspect-template-%s" , uuid .New ().String ())
87+ deidentifyTemplateID := fmt .Sprintf ("model-armor-deidentify-template-%s" , uuid .New ().String ())
88+ apiEndpoint := fmt .Sprintf ("dlp.%s.rep.googleapis.com:443" , locationID )
89+ parent := fmt .Sprintf ("projects/%s/locations/%s" , projectID , locationID )
90+
91+ infoTypes := []* dlppb.InfoType {
92+ {Name : "EMAIL_ADDRESS" },
93+ {Name : "PHONE_NUMBER" },
94+ {Name : "US_INDIVIDUAL_TAXPAYER_IDENTIFICATION_NUMBER" },
95+ }
96+
97+ ctx := context .Background ()
98+ dlpClient , err := dlp .NewClient (ctx , option .WithEndpoint (apiEndpoint ))
99+ if err != nil {
100+ t .Fatalf ("Getting error while creating the client: %v" , err )
101+ }
102+ defer dlpClient .Close ()
103+
104+ inspectRequest := & dlppb.CreateInspectTemplateRequest {
105+ Parent : parent ,
106+ TemplateId : inspectTemplateID ,
107+ InspectTemplate : & dlppb.InspectTemplate {
108+ InspectConfig : & dlppb.InspectConfig {
109+ InfoTypes : infoTypes ,
110+ },
111+ },
112+ }
113+ inspectResponse , err := dlpClient .CreateInspectTemplate (ctx , inspectRequest )
114+ if err != nil {
115+ t .Fatal (err )
116+ }
117+
118+ deidentifyRequest := & dlppb.CreateDeidentifyTemplateRequest {
119+ Parent : parent ,
120+ TemplateId : deidentifyTemplateID ,
121+ DeidentifyTemplate : & dlppb.DeidentifyTemplate {
122+ DeidentifyConfig : & dlppb.DeidentifyConfig {
123+ Transformation : & dlppb.DeidentifyConfig_InfoTypeTransformations {
124+ InfoTypeTransformations : & dlppb.InfoTypeTransformations {
125+ Transformations : []* dlppb.InfoTypeTransformations_InfoTypeTransformation {
126+ {
127+ InfoTypes : []* dlppb.InfoType {},
128+ PrimitiveTransformation : & dlppb.PrimitiveTransformation {
129+ Transformation : & dlppb.PrimitiveTransformation_ReplaceConfig {
130+ ReplaceConfig : & dlppb.ReplaceValueConfig {
131+ NewValue : & dlppb.Value {
132+ Type : & dlppb.Value_StringValue {StringValue : "REDACTED" },
133+ },
134+ },
135+ },
136+ },
137+ },
138+ },
139+ },
140+ },
141+ },
142+ },
143+ }
144+ deidentifyResponse , err := dlpClient .CreateDeidentifyTemplate (ctx , deidentifyRequest )
145+ if err != nil {
146+ t .Fatal (err )
147+ }
148+
149+ // Cleanup the templates after test.
150+ defer func () {
151+ time .Sleep (5 * time .Second )
152+ err := dlpClient .DeleteInspectTemplate (ctx , & dlppb.DeleteInspectTemplateRequest {Name : inspectResponse .Name })
153+ if err != nil {
154+ t .Errorf ("failed to delete inspect template: %v, err: %v" , inspectResponse .Name , err )
155+ }
156+ err = dlpClient .DeleteDeidentifyTemplate (ctx , & dlppb.DeleteDeidentifyTemplateRequest {Name : deidentifyResponse .Name })
157+ if err != nil {
158+ t .Errorf ("failed to delete deidentify template: %v, err: %v" , deidentifyResponse .Name , err )
159+ }
160+ }()
161+
162+ return inspectResponse .Name , deidentifyResponse .Name
163+ }
164+
165+ // TestCreateModelArmorTemplateWithAdvancedSDP tests creating a
166+ // Model Armor template with advanced SDP using DLP templates.
167+ func TestCreateModelArmorTemplateWithAdvancedSDP (t * testing.T ) {
168+ tc := testutil .SystemTest (t )
169+
170+ templateID := fmt .Sprintf ("test-model-armor-%s" , uuid .New ().String ())
171+ inspectTemplateName , deideintifyTemplateName := testSDPTemplate (t , tc .ProjectID , testLocation (t ))
172+ templateName := fmt .Sprintf ("projects/%s/locations/%s/templates/%s" , tc .ProjectID , testLocation (t ), templateID )
173+ var buf bytes.Buffer
174+ if err := createModelArmorTemplateWithAdvancedSDP (& buf , tc .ProjectID , testLocation (t ), templateID , inspectTemplateName , deideintifyTemplateName ); err != nil {
175+ t .Fatal (err )
176+ }
177+ defer testCleanupTemplate (t , templateName )
178+
179+ if got , want := buf .String (), "Created Template with advanced SDP: " ; ! strings .Contains (got , want ) {
180+ t .Errorf ("createModelArmorTemplateWithAdvancedSDP: expected %q to contain %q" , got , want )
181+ }
182+ }
183+
81184// TestCreateModelArmorTemplate verifies the creation of a Model Armor template.
82185// It ensures the output contains a confirmation message after creation.
83186func TestCreateModelArmorTemplate (t * testing.T ) {
0 commit comments