feat(modelarmor): code refactor

Author: tirthrajsinh-zala-crest

modelarmor/create_template.go

@@ -6,7 +6,6 @@ require (
 	cloud.google.com/go/modelarmor v0.1.0
 	github.com/GoogleCloudPlatform/golang-samples v0.0.0-20250404170905-0aca11152736
 	github.com/google/uuid v1.6.0
-	github.com/joho/godotenv v1.5.1
 	google.golang.org/api v0.226.0
 	google.golang.org/grpc v1.71.0
 )

modelarmor/go.mod

@@ -67,8 +67,6 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.5 h1:VgzTY2jogw3xt39CusE
 github.com/googleapis/enterprise-certificate-proxy v0.3.5/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA=
 github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q=
 github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
-github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
-github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
 github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
 github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

modelarmor/go.sum

@@ -27,21 +27,16 @@ import (
 
 	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
 	"github.com/google/uuid"
-	"github.com/joho/godotenv"
 	"google.golang.org/api/option"
 	grpccodes "google.golang.org/grpc/codes"
 	grpcstatus "google.golang.org/grpc/status"
 )
 
+// testLocation returns the test location from the environment variable `GOLANG_SAMPLES_LOCATION`.
+// If not set, the test is skipped.
 func testLocation(t *testing.T) string {
 	t.Helper()
 
-	// Load the test.env file
-	err := godotenv.Load("./testdata/env/test.env")
-	if err != nil {
-		t.Fatalf("failed to load test environment file: %v", err)
-	}
-
 	v := os.Getenv("GOLANG_SAMPLES_LOCATION")
 	if v == "" {
 		t.Skip("testIamUser: missing GOLANG_SAMPLES_LOCATION")
@@ -50,36 +45,76 @@ func testLocation(t *testing.T) string {
 	return v
 }
 
+// testClient initializes and returns a new Model Armor client and context
+// for use in tests. The client is created using the location-specific endpoint.
 func testClient(t *testing.T) (*modelarmor.Client, context.Context) {
 	t.Helper()
 
 	ctx := context.Background()
 
-	locationId := testLocation(t)
+	locationID := testLocation(t)
 
-	//Endpoint to send the request to regional server
-	client, err := modelarmor.NewClient(ctx,
-		option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationId)),
-	)
+	// Create options for Model Armor client.
+	opts := option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID))
+	// Create the Model Armor client.
+	client, err := modelarmor.NewClient(ctx, opts)
 	if err != nil {
 		t.Fatalf("failed to create client: %v", err)
 	}
 
 	return client, ctx
 }
 
+// testModelArmorTemplate creates a new ModelArmor template for use in tests.
+// It returns the created template or an error.
+func testModelArmorTemplate(t *testing.T, templateID string) (*modelarmorpb.Template, error) {
+	t.Helper()
+	tc := testutil.SystemTest(t)
+	locationID := testLocation(t)
+	client, ctx := testClient(t)
+
+	template := &modelarmorpb.Template{
+		FilterConfig: &modelarmorpb.FilterConfig{
+			PiAndJailbreakFilterSettings: &modelarmorpb.PiAndJailbreakFilterSettings{
+				FilterEnforcement: modelarmorpb.PiAndJailbreakFilterSettings_ENABLED,
+				ConfidenceLevel:   modelarmorpb.DetectionConfidenceLevel_MEDIUM_AND_ABOVE,
+			},
+			MaliciousUriFilterSettings: &modelarmorpb.MaliciousUriFilterSettings{
+				FilterEnforcement: modelarmorpb.MaliciousUriFilterSettings_ENABLED,
+			},
+		},
+	}
+
+	req := &modelarmorpb.CreateTemplateRequest{
+		Parent:     fmt.Sprintf("projects/%s/locations/%s", tc.ProjectID, locationID),
+		TemplateId: templateID,
+		Template:   template,
+	}
+
+	response, err := client.CreateTemplate(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create template: %v", err)
+	}
+
+	return response, err
+}
+
+// testCleanupTemplate deletes a Model Armor template by name, used for cleanup
+// after tests. Ignores errors if the template is already deleted.
 func testCleanupTemplate(t *testing.T, templateName string) {
 	t.Helper()
 
 	client, ctx := testClient(t)
 	if err := client.DeleteTemplate(ctx, &modelarmorpb.DeleteTemplateRequest{Name: templateName}); err != nil {
 		if terr, ok := grpcstatus.FromError(err); !ok || terr.Code() != grpccodes.NotFound {
-			t.Fatalf("testCleanupTemplate: failed to delete template %s: %v", templateName, err)
+			t.Fatalf("testCleanupTemplate: failed to delete template: %v", err)
 		}
 	}
 
 }
 
+// TestScreenPDFFile scrrens the pdf file content and Sanitize 
+// the content with the Model Armor.
 func TestScreenPDFFile(t *testing.T) {
 	pdfContentBase := `"JVBERi0xLjQKJdPr6eEKMSAwIG9iago8PC9UaXRsZSAoVW50aXRsZWQgZG9jdW1lbnQpCi9Qcm9kdWNlciAoU2tp
 YS9QREYgbTEzNSBHb29nbGUgRG9jcyBSZW5kZXJlcik+PgplbmRvYmoKMyAwIG9iago8PC9jYSAxCi9CTSAvTm9y
@@ -391,18 +426,19 @@ MDAwMTg5MTQgMDAwMDAgbiAKMDAwMDAxOTMyOSAwMDAwMCBuIAp0cmFpbGVyCjw8L1NpemUgMTkKL1Jv
 MCBSCi9JbmZvIDEgMCBSPj4Kc3RhcnR4cmVmCjE5ODQ0CiUlRU9GCg=="`
 	tc := testutil.SystemTest(t)
 	templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
-
-	var b bytes.Buffer
-	if _, err := createModelArmorTemplate(&b, tc.ProjectID, testLocation(t), templateID); err != nil {
+	locationID := testLocation(t)
+	templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)
+	var buf bytes.Buffer
+	if _, err := testModelArmorTemplate(t, templateID); err != nil {
 		t.Fatal(err)
 	}
-	defer testCleanupTemplate(t, fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, testLocation(t), templateID))
+	defer testCleanupTemplate(t, templateName)
 
-	if _, err := screenPDFFile(&b, tc.ProjectID, testLocation(t), templateID, pdfContentBase); err != nil {
+	if err := screenPDFFile(&buf, tc.ProjectID, testLocation(t), templateID, pdfContentBase); err != nil {
 		t.Fatal(err)
 	}
 
-	if got, want := b.String(), "PDF screening sanitization result: "; !strings.Contains(got, want) {
+	if got, want := buf.String(), "PDF screening sanitization result: "; !strings.Contains(got, want) {
 		t.Errorf("screenPdf: expected %q to contain %q", got, want)
 	}
 }

modelarmor/modelarmor_test.go

@@ -32,41 +32,20 @@ import (
 //
 // This method screens a PDF file based on the project, location, and template settings.
 //
-// Args:
-//
-//	w io.Writer: The writer to use for logging.
-//	projectID string: The ID of the project.
-//	locationID string: The ID of the location.
-//	templateID string: The ID of the template.
-//	pdfContentBase64 string: The base64-encoded content of the PDF file.
-//
-// Returns:
-//
-//	*modelarmorpb.SanitizeUserPromptResponse: The response from screening the PDF file.
-//	error: Any error that occurred during screening.
-//
-// Example:
-//
-//	response, err := screenPDFFile(
-//	    os.Stdout,
-//	    "my-project",
-//	    "my-location",
-//	    "my-template",
-//	    "base64-encoded-pdf-content",
-//	)
-//	if err != nil {
-//	    log.Fatal(err)
-//	}
-//	fmt.Println(response)
-func screenPDFFile(w io.Writer, projectID, locationID, templateID, pdfContentBase64 string) (*modelarmorpb.SanitizeUserPromptResponse, error) {
+// w io.Writer: The writer to use for logging.
+// projectID string: The ID of the project.
+// locationID string: The ID of the location.
+// templateID string: The ID of the template.
+// pdfContentBase64 string: The base64-encoded content of the PDF file.
+func screenPDFFile(w io.Writer, projectID, locationID, templateID, pdfContentBase64 string) error {
 	ctx := context.Background()
 
+	// Create options for Model Armor client.
+	opts := option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID))
 	// Create the Model Armor client.
-	client, err := modelarmor.NewClient(ctx,
-		option.WithEndpoint(fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID)),
-	)
+	client, err := modelarmor.NewClient(ctx, opts)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create client: %v", err)
+		return fmt.Errorf("failed to create client: %w", err)
 	}
 	defer client.Close()
 
@@ -89,13 +68,13 @@ func screenPDFFile(w io.Writer, projectID, locationID, templateID, pdfContentBas
 	// Sanitize the user prompt.
 	response, err := client.SanitizeUserPrompt(ctx, req)
 	if err != nil {
-		return nil, fmt.Errorf("failed to sanitize PDF content for template %s: %v", templateID, err)
+		return fmt.Errorf("failed to sanitize PDF content for template %s: %w", templateID, err)
 	}
 
 	// Sanitization Result.
 	fmt.Fprintf(w, "PDF screening sanitization result: %v\n", response)
 
-	// [END modelarmor_screen_pdf_file]
-
-	return response, nil
+	return err
 }
+
+// [END modelarmor_screen_pdf_file]