diff --git a/packages/synthetics-sdk-broken-links/src/storage_func.ts b/packages/synthetics-sdk-broken-links/src/storage_func.ts index d0ed7cc5..e906ac29 100644 --- a/packages/synthetics-sdk-broken-links/src/storage_func.ts +++ b/packages/synthetics-sdk-broken-links/src/storage_func.ts @@ -12,14 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +import * as path from 'path'; import { Storage, Bucket } from '@google-cloud/storage'; import { BaseError, - BrokenLinksResultV1_BrokenLinkCheckerOptions_ScreenshotOptions_ScreenshotCondition, + BrokenLinksResultV1_BrokenLinkCheckerOptions, + BrokenLinksResultV1_BrokenLinkCheckerOptions_ScreenshotOptions_ScreenshotCondition as ApiScreenshotCondition, resolveProjectId, getExecutionRegion, + BrokenLinksResultV1_SyntheticLinkResult_ScreenshotOutput as ApiScreenshotOutput, } from '@google-cloud/synthetics-sdk-api'; +export interface StorageParameters { + storageClient: Storage | null; + bucket: Bucket | null; + uptimeId: string; + executionId: string; +} + /** * Attempts to get an existing storage bucket if provided by the user OR * create/use a dedicated synthetics bucket. @@ -97,13 +107,9 @@ export async function getOrCreateStorageBucket( */ export function createStorageClientIfStorageSelected( errors: BaseError[], - storage_condition: BrokenLinksResultV1_BrokenLinkCheckerOptions_ScreenshotOptions_ScreenshotCondition + storageCondition: ApiScreenshotCondition ): Storage | null { - if ( - storage_condition === - BrokenLinksResultV1_BrokenLinkCheckerOptions_ScreenshotOptions_ScreenshotCondition.NONE - ) - return null; + if (storageCondition === ApiScreenshotCondition.NONE) return null; try { return new Storage(); @@ -117,3 +123,77 @@ export function createStorageClientIfStorageSelected( return null; } } + +/** + * Uploads a screenshot to Google Cloud Storage. + * + * @param screenshot - Base64-encoded screenshot data. + * @param filename - Desired filename for the screenshot. + * @param storageParams - An object containing storageClient and bucket. + * @param options - Broken links checker options. + * @returns An ApiScreenshotOutput object indicating success or a screenshot_error. + */ +export async function uploadScreenshotToGCS( + screenshot: string, + filename: string, + storageParams: StorageParameters, + options: BrokenLinksResultV1_BrokenLinkCheckerOptions +): Promise { + const screenshot_output: ApiScreenshotOutput = { + screenshot_file: '', + screenshot_error: {} as BaseError, + }; + try { + // Early exit if storage is not properly configured + if (!storageParams.storageClient || !storageParams.bucket) { + return screenshot_output; + } + + // Construct the destination path within the bucket if given + let writeDestination = options.screenshot_options!.storage_location + ? getFolderNameFromStorageLocation( + options.screenshot_options!.storage_location + ) + : ''; + + // Ensure writeDestination ends with a slash for proper path joining + if (writeDestination && !writeDestination.endsWith('/')) { + writeDestination += '/'; + } + + writeDestination = path.join( + writeDestination, + storageParams.uptimeId, + storageParams.executionId, + filename + ); + + // Upload to GCS + await storageParams.bucket.file(writeDestination).save(screenshot, { + contentType: 'image/png', + }); + + screenshot_output.screenshot_file = writeDestination; + } catch (err) { + // Handle upload errors + if (err instanceof Error) process.stderr.write(err.message); + screenshot_output.screenshot_error = { + error_type: 'StorageFileUploadError', + error_message: `Failed to upload screenshot for ${filename}. Please reference server logs for further information.`, + }; + } + + return screenshot_output; +} + +// Helper function to extract folder name for a given storage location. If there +// is no '/' present then the storageLocation is just a folder +export function getFolderNameFromStorageLocation( + storageLocation: string +): string { + const firstSlashIndex = storageLocation.indexOf('/'); + if (firstSlashIndex === -1) { + return ''; + } + return storageLocation.substring(firstSlashIndex + 1); +} diff --git a/packages/synthetics-sdk-broken-links/test/unit/storage_func.spec.ts b/packages/synthetics-sdk-broken-links/test/unit/storage_func.spec.ts index 3842fbee..969fec43 100644 --- a/packages/synthetics-sdk-broken-links/test/unit/storage_func.spec.ts +++ b/packages/synthetics-sdk-broken-links/test/unit/storage_func.spec.ts @@ -14,12 +14,16 @@ import { expect } from 'chai'; import sinon from 'sinon'; -import { Storage, Bucket } from '@google-cloud/storage'; +import { Storage, Bucket, File } from '@google-cloud/storage'; import * as sdkApi from '@google-cloud/synthetics-sdk-api'; import { createStorageClientIfStorageSelected, + getFolderNameFromStorageLocation, getOrCreateStorageBucket, + StorageParameters, + uploadScreenshotToGCS, } from '../../src/storage_func'; +import { BrokenLinksResultV1_BrokenLinkCheckerOptions } from '@google-cloud/synthetics-sdk-api'; const proxyquire = require('proxyquire'); // global test vars @@ -33,6 +37,7 @@ describe('GCM Synthetics Broken Links storage_func suite testing', () => { '@google-cloud/synthetics-sdk-api': { getExecutionRegion: () => 'test-region', resolveProjectId: () => 'test-project-id', + getOrCreateStorageBucket: () => getOrCreateStorageBucket, }, }); @@ -94,10 +99,14 @@ describe('GCM Synthetics Broken Links storage_func suite testing', () => { const result = await storageFunc.getOrCreateStorageBucket( storageClientStub, - '', + TEST_BUCKET_NAME + '/fake-folder', [] ); expect(result).to.equal(bucketStub); + sinon.assert.calledWithExactly( + storageClientStub.bucket, + TEST_BUCKET_NAME + ); sinon.assert.notCalled(bucketStub.create); }); @@ -148,4 +157,157 @@ describe('GCM Synthetics Broken Links storage_func suite testing', () => { expect(result).to.be.an.instanceOf(Storage); }); }); + + describe('uploadScreenshotToGCS', () => { + let storageClientStub: sinon.SinonStubbedInstance; + let bucketStub: sinon.SinonStubbedInstance; + + const screenshotData = 'encoded-image-data'; + const filename = 'test-screenshot.png'; + + beforeEach(() => { + storageClientStub = sinon.createStubInstance(Storage); + bucketStub = sinon.createStubInstance(Bucket); + storageClientStub.bucket.returns(bucketStub); + }); + + afterEach(() => { + sinon.restore(); + }); + + it('should upload the screenshot and return the write_destination', async () => { + const storageParams = { + storageClient: storageClientStub, + bucket: bucketStub, + uptimeId: 'uptime123', + executionId: 'exec456', + }; + const options = { + screenshot_options: { storage_location: 'bucket/folder1/folder2' }, + } as BrokenLinksResultV1_BrokenLinkCheckerOptions; + const expectedWriteDestination = + 'folder1/folder2/uptime123/exec456/test-screenshot.png'; + + const successPartialFileMock: Partial = { + save: sinon.stub().resolves(), + }; + bucketStub.file.returns(successPartialFileMock as File); + + const result = await uploadScreenshotToGCS( + screenshotData, + filename, + storageParams, + options + ); + + expect(result.screenshot_file).to.equal(expectedWriteDestination); + expect(result.screenshot_error).to.deep.equal({}); + }); + + it('should handle GCS upload errors', async () => { + const storageParams: StorageParameters = { + storageClient: storageClientStub, + bucket: bucketStub, + uptimeId: '', + executionId: '', + }; + const options = { + screenshot_options: {}, + } as BrokenLinksResultV1_BrokenLinkCheckerOptions; + + const gcsError = new Error('Simulated GCS upload error'); + const failingPartialFileMock: Partial = { + save: sinon.stub().throws(gcsError), + }; + bucketStub.file.returns(failingPartialFileMock as File); + + const result = await uploadScreenshotToGCS( + screenshotData, + filename, + storageParams, + options + ); + + expect(result.screenshot_file).to.equal(''); + expect(result.screenshot_error).to.deep.equal({ + error_type: 'StorageFileUploadError', + error_message: `Failed to upload screenshot for ${filename}. Please reference server logs for further information.`, + }); + }); + + describe('Invalid Storage Configuration', () => { + const emptyScreenshotData = ''; + const emptyFilename = ''; + const emptyOptions = {} as BrokenLinksResultV1_BrokenLinkCheckerOptions; + it('should return an empty result if storageClient is null', async () => { + // Missing storageClient + const storageParams = { + storageClient: null, + bucket: bucketStub, + uptimeId: '', + executionId: '', + }; + + const result = await uploadScreenshotToGCS( + emptyScreenshotData, + emptyFilename, + storageParams, + emptyOptions + ); + + expect(result).to.deep.equal({ + screenshot_file: '', + screenshot_error: {}, + }); + }); + + it('should return an empty result if bucket is null', async () => { + // Missing bucket + const storageParams = { + storageClient: storageClientStub, + bucket: null, + uptimeId: '', + executionId: '', + }; + + const result = await uploadScreenshotToGCS( + emptyScreenshotData, + emptyFilename, + storageParams, + emptyOptions + ); + + expect(result).to.deep.equal({ + screenshot_file: '', + screenshot_error: {}, + }); + }); + }); + }); + + describe('getFolderNameFromStorageLocation', () => { + it('should extract folder name when storage location has a slash', () => { + const storageLocation = 'some-bucket/folder1/folder2'; + const expectedFolderName = 'folder1/folder2'; + + const result = getFolderNameFromStorageLocation(storageLocation); + expect(result).to.equal(expectedFolderName); + }); + + it('should return an empty string if storage location has no slash', () => { + const storageLocation = 'my-bucket'; + const expectedFolderName = ''; + + const result = getFolderNameFromStorageLocation(storageLocation); + expect(result).to.equal(expectedFolderName); + }); + + it('should return an empty string if given an empty string', () => { + const storageLocation = ''; + const expectedFolderName = ''; + + const result = getFolderNameFromStorageLocation(storageLocation); + expect(result).to.equal(expectedFolderName); + }); + }); });