take-screenshots (#107)

Author: adamfweidman

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<ApiScreenshotOutput> {
+  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);
+}

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<Storage>;
+    let bucketStub: sinon.SinonStubbedInstance<Bucket>;
+
+    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<File> = {
+        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<File> = {
+        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);
+    });
+  });
 });