fixup

components/consumers/pdf/main_test.go

@@ -0,0 +1,62 @@
+// Package main of the pdf consumer implements a simple consumer for
+// applying a go-template to a smithy scan, converting the result to pdf and then
+// uploading the result to the S3 bucket passed as an argument
+// the consumer expects the environment variables
+// AWS_ACCESS_KEY_ID
+// AWS_SECRET_ACCESS_KEY
+// to be set along with the "bucket" and "region" arguments to be passed
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	v1 "github.com/smithy-security/smithy/api/proto/v1"
+
+	playwright "github.com/smithy-security/smithy/pkg/playwright/mock"
+	s3mock "github.com/smithy-security/smithy/pkg/s3/mock"
+	"github.com/smithy-security/smithy/pkg/testdata"
+)
+
+func Test_run(t *testing.T) {
+	mockClient, err := playwright.NewMockClient()
+	require.NoError(t, err)
+
+	pdfCalled := false
+	expected := []byte("this is a pdf")
+	mockClient.GetPDFOfPageCallBack = func(s1, s2 string) ([]byte, error) {
+		pdfCalled = true
+		return expected, nil
+	}
+
+	mockS3Client, err := s3mock.NewMockClient("")
+	require.NoError(t, err)
+	s3Called := false
+	mockS3Client.UpsertCallback = func(s1, s2 string, b []byte) error {
+		s3Called = true
+		return nil
+	}
+
+	err = run([]v1.EnrichedLaunchToolResponse{testdata.EnrichedLaunchToolResponse}, mockClient, mockS3Client)
+	require.NoError(t, err)
+	require.True(t, pdfCalled)
+	require.True(t, s3Called)
+
+}
+
+func Test_buildPdf(t *testing.T) {
+	mockClient, err := playwright.NewMockClient()
+	require.NoError(t, err)
+
+	called := false
+	expected := []byte("this is a pdf")
+	mockClient.GetPDFOfPageCallBack = func(s1, s2 string) ([]byte, error) {
+		called = true
+		return expected, nil
+	}
+	_, result, err := buildPdf([]v1.EnrichedLaunchToolResponse{testdata.EnrichedLaunchToolResponse}, mockClient)
+	require.NoError(t, err)
+	require.Equal(t, called, true)
+	require.Equal(t, result, expected)
+}

pkg/playwright/mock/wrapper.go

@@ -0,0 +1,19 @@
+package mock
+
+type MockClient struct {
+	StopCallBack         func() error
+	GetPDFOfPageCallBack func(string, string) ([]byte, error)
+}
+
+// NewClient returns an actual github client
+func NewMockClient() (MockClient, error) {
+	return MockClient{}, nil
+}
+
+func (c MockClient) Stop() error {
+	return c.StopCallBack()
+}
+
+func (c MockClient) GetPDFOfPage(page, storePath string) ([]byte, error) {
+	return c.GetPDFOfPageCallBack(page, storePath)
+}

pkg/playwright/wrapper.go

@@ -0,0 +1,58 @@
+package wrapper
+
+import (
+	"github.com/go-errors/errors"
+	"github.com/playwright-community/playwright-go"
+)
+
+// Wrapper is the wrapper interface that allows the playwright client to be pluggable
+type Wrapper interface {
+	Stop() error
+	GetPDFOfPage(string, string) ([]byte, error)
+}
+
+// Client is the wrapper around google's go-github client
+type Client struct {
+	playwright *playwright.Playwright
+}
+
+// NewClient returns an actual github client
+func NewClient() (Client, error) {
+	pw, err := playwright.Run()
+	if err != nil {
+		return Client{}, err
+	}
+	// create new playwright client
+	return Client{
+		playwright: pw,
+	}, nil
+}
+
+func (c Client) Stop() error {
+	return c.playwright.Stop()
+}
+
+func (c Client) GetPDFOfPage(page, storePath string) ([]byte, error) {
+	browser, err := c.playwright.Chromium.Launch()
+	if err != nil {
+		return nil, err
+	}
+
+	currentContext, err := browser.NewContext()
+	if err != nil {
+		return nil, err
+	}
+
+	newPage, err := currentContext.NewPage()
+	if err != nil {
+		return nil, err
+	}
+
+	if _, err = newPage.Goto(page); err != nil {
+		return nil, errors.Errorf("could not goto page %s in the browser: %w", page, err)
+	}
+
+	return newPage.PDF(playwright.PagePdfOptions{
+		Path: playwright.String(storePath),
+	})
+}

pkg/s3/mock/wrapper.go

@@ -0,0 +1,16 @@
+package wrapper
+
+// Client is the wrapper around s3 sdk
+type Client struct {
+	UpsertCallback func(string, string, []byte) error
+}
+
+// NewClient returns a client
+func NewMockClient(region string) (Client, error) {
+	// create new playwright client
+	return Client{}, nil
+}
+
+func (c Client) UpsertFile(filename, bucket string, pdfBytes []byte) error {
+	return c.UpsertCallback(filename, bucket, pdfBytes)
+}

pkg/s3/wrapper.go

@@ -0,0 +1,65 @@
+package wrapper
+
+import (
+	"bytes"
+	"log/slog"
+	"os"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+	"github.com/go-errors/errors"
+)
+
+// Wrapper is the wrapper interface that allows the playwright client to be pluggable
+type Wrapper interface {
+	// UpsertFile inserts or replaces a file in s3 if it doesn't exist
+	UpsertFile(string, string, []byte) error
+}
+
+// Client is the wrapper around google's go-github client
+type Client struct {
+	session *session.Session
+}
+
+// NewClient returns a client
+func NewClient(region string) (Client, error) {
+	sess, err := session.NewSession(&aws.Config{Region: aws.String(region)})
+	if err != nil {
+		return Client{}, errors.Errorf("unable to start session with AWS API: %w", err)
+	}
+	// create new playwright client
+	return Client{
+		session: sess,
+	}, nil
+}
+
+func (c Client) UpsertFile(filename, bucket string, pdfBytes []byte) error {
+	//#nosec:G304
+	data, err := os.ReadFile(filename) //#nosec:G304
+	if err != nil {
+		return errors.Errorf("could not open file: %w", err)
+	}
+
+	uploader := s3manager.NewUploader(c.session)
+	_, err = uploader.Upload(&s3manager.UploadInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(filename),
+		Body:   bytes.NewReader(data),
+	})
+	if err != nil {
+		return errors.Errorf("unable to upload %s to %s: %w", filename, bucket, err)
+	}
+
+	_, err = uploader.Upload(&s3manager.UploadInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(filename),
+		Body:   bytes.NewReader(pdfBytes),
+	})
+	if err != nil {
+		return errors.Errorf("unable to upload %s to %s: %w", filename, bucket, err)
+	}
+
+	slog.Info("uploaded", "filename", filename, "to", "bucket", bucket, "successfully")
+	return nil
+}