Incorporated review comments.

Author: suratdas

src/compare/compare.module.ts

@@ -3,9 +3,24 @@ import { CompareService } from './compare.service';
 import { LookSameService } from './libs/looks-same/looks-same.service';
 import { OdiffService } from './libs/odiff/odiff.service';
 import { PixelmatchService } from './libs/pixelmatch/pixelmatch.service';
+import { AWSS3Service } from 'src/shared/static/aws-s3.servce.';
+import { HardDiskService } from 'src/shared/static/hard-disk.service';
+import { STATIC_SERVICE, StaticService } from 'src/shared/static/static-service.interface';
 
 @Module({
-  providers: [CompareService, PixelmatchService, LookSameService, OdiffService],
+  providers: [
+    {
+      provide: STATIC_SERVICE,
+      useFactory: (): StaticService => {
+        const isAWSDefined = process.env.USE_AWS_S3_BUCKET?.trim().toLowerCase() === 'true';
+        return isAWSDefined ? new AWSS3Service() : new HardDiskService();
+      },
+    },
+    CompareService,
+    PixelmatchService,
+    LookSameService,
+    OdiffService,
+  ],
   exports: [CompareService],
 })
 export class CompareModule {}

src/compare/libs/looks-same/looks-same.service.ts

@@ -1,7 +1,7 @@
-import { Injectable, Logger } from '@nestjs/common';
+import { Inject, Injectable, Logger } from '@nestjs/common';
 import { TestStatus } from '@prisma/client';
 import { PNG } from 'pngjs';
-import { StaticService } from '../../../shared/static/static.service';
+import { STATIC_SERVICE, StaticService } from '../../../shared/static/static-service.interface';
 import { DiffResult } from '../../../test-runs/diffResult';
 import { applyIgnoreAreas, parseConfig } from '../../utils';
 import { ImageComparator } from '../image-comparator.interface';
@@ -21,7 +21,7 @@ export const DEFAULT_CONFIG: LooksSameConfig = {
 export class LookSameService implements ImageComparator {
   private readonly logger: Logger = new Logger(LookSameService.name);
 
-  constructor(private staticService: StaticService) {}
+  constructor(@Inject(STATIC_SERVICE) private staticService: StaticService) {}
 
   parseConfig(configJson: string): LooksSameConfig {
     return parseConfig(configJson, DEFAULT_CONFIG, this.logger);

src/compare/libs/odiff/odiff.service.ts

@@ -1,6 +1,6 @@
-import { Injectable, Logger } from '@nestjs/common';
+import { Inject, Injectable, Logger } from '@nestjs/common';
 import { TestStatus } from '@prisma/client';
-import { StaticService } from '../../../shared/static/static.service';
+import { STATIC_SERVICE, StaticService } from '../../../shared/static/static-service.interface';
 import { DiffResult } from '../../../test-runs/diffResult';
 import { parseConfig } from '../../utils';
 import { ImageComparator } from '../image-comparator.interface';
@@ -21,7 +21,7 @@ export const DEFAULT_CONFIG: OdiffConfig = {
 export class OdiffService implements ImageComparator {
   private readonly logger: Logger = new Logger(OdiffService.name);
 
-  constructor(private staticService: StaticService) {}
+  constructor(@Inject(STATIC_SERVICE) private staticService: StaticService) {}
 
   parseConfig(configJson: string): OdiffConfig {
     return parseConfig(configJson, DEFAULT_CONFIG, this.logger);

src/compare/libs/pixelmatch/pixelmatch.service.ts

@@ -1,8 +1,8 @@
-import { Injectable, Logger } from '@nestjs/common';
+import { Inject, Injectable, Logger } from '@nestjs/common';
 import { TestStatus } from '@prisma/client';
 import Pixelmatch from 'pixelmatch';
 import { PNG } from 'pngjs';
-import { StaticService } from '../../../shared/static/static.service';
+import { STATIC_SERVICE, StaticService } from '../../../shared/static/static-service.interface';
 import { DiffResult } from '../../../test-runs/diffResult';
 import { scaleImageToSize, applyIgnoreAreas, parseConfig } from '../../utils';
 import { DIFF_DIMENSION_RESULT, EQUAL_RESULT, NO_BASELINE_RESULT } from '../consts';
@@ -16,7 +16,7 @@ export const DEFAULT_CONFIG: PixelmatchConfig = { threshold: 0.1, ignoreAntialia
 export class PixelmatchService implements ImageComparator {
   private readonly logger: Logger = new Logger(PixelmatchService.name);
 
-  constructor(private staticService: StaticService) {}
+  constructor(@Inject(STATIC_SERVICE) private staticService: StaticService) {}
 
   parseConfig(configJson: string): PixelmatchConfig {
     return parseConfig(configJson, DEFAULT_CONFIG, this.logger);

src/main.ts

@@ -6,8 +6,8 @@ import { join } from 'path';
 import * as bodyParser from 'body-parser';
 import { readFileSync, existsSync } from 'fs';
 import { HttpsOptions } from '@nestjs/common/interfaces/external/https-options.interface';
-import { IMAGE_PATH } from './shared/static/static.service';
 import { NestExpressApplication } from '@nestjs/platform-express';
+import { IMAGE_PATH } from './shared/static/common-file-service';
 
 function getHttpsOptions(): HttpsOptions | null {
   const keyPath = './secrets/ssl.key';

src/shared/shared.module.ts

@@ -1,15 +1,30 @@
 import { forwardRef, Global, Module } from '@nestjs/common';
-import { StaticService } from './static/static.service';
+import { STATIC_SERVICE, StaticService } from './static/static-service.interface';
 import { EventsGateway } from '../shared/events/events.gateway';
 import { PrismaService } from '../prisma/prisma.service';
 import { TasksService } from './tasks/tasks.service';
 import { TestVariationsModule } from '../test-variations/test-variations.module';
 import { StaticController } from './static/static.controller';
+import { AWSS3Service } from './static/aws-s3.servce.';
+import { HardDiskService } from './static/hard-disk.service';
 
 @Global()
 @Module({
-  providers: [StaticService, EventsGateway, PrismaService, TasksService],
-  exports: [StaticService, EventsGateway, PrismaService],
+  providers: [
+    {
+      provide: STATIC_SERVICE,
+      useFactory: (): StaticService => {
+        const isAWSDefined = process.env.USE_AWS_S3_BUCKET?.trim().toLowerCase() === 'true';
+        return isAWSDefined ? new AWSS3Service() : new HardDiskService();
+      },
+    },
+    AWSS3Service,
+    HardDiskService,
+    EventsGateway,
+    PrismaService,
+    TasksService,
+  ],
+  exports: [STATIC_SERVICE, EventsGateway, PrismaService],
   imports: [forwardRef(() => TestVariationsModule)],
   controllers: [StaticController],
 })

src/shared/static/aws-s3.servce..ts

@@ -0,0 +1,94 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { readFileSync, createWriteStream } from 'fs';
+import { PNG, PNGWithMetadata } from 'pngjs';
+
+import { S3Client, PutObjectCommand, DeleteObjectCommand, GetObjectCommand } from '@aws-sdk/client-s3';
+import { Readable } from 'stream';
+import { StaticService } from './static-service.interface';
+import { CommonFileService } from './common-file-service';
+
+@Injectable()
+export class AWSS3Service extends CommonFileService implements StaticService {
+  private readonly AWS_ACCESS_KEY_ID = process.env.AWS_ACCESS_KEY_ID;
+  private readonly AWS_SECRET_ACCESS_KEY = process.env.AWS_SECRET_ACCESS_KEY;
+  private readonly AWS_REGION = process.env.AWS_REGION;
+  private readonly AWS_S3_BUCKET_NAME = process.env.AWS_S3_BUCKET_NAME;
+
+  private s3Client: S3Client;
+
+  constructor() {
+    super();
+    this.s3Client = new S3Client({
+      credentials: {
+        accessKeyId: this.AWS_ACCESS_KEY_ID,
+        secretAccessKey: this.AWS_SECRET_ACCESS_KEY,
+      },
+      region: this.AWS_REGION,
+    });
+    this.logger = new Logger(AWSS3Service.name);
+    this.logger.log('AWS S3 service is being used for file storage.');
+    setInterval(() => this.cleanupQueuedFiles(), this.DELETE_INTERVAL);
+    this.cleanUpAllFiles();
+  }
+
+  async saveImage(type: 'screenshot' | 'diff' | 'baseline', imageBuffer: Buffer): Promise<string> {
+    const { imageName } = this.generateNewImage(type);
+    try {
+      await this.s3Client.send(
+        new PutObjectCommand({
+          Bucket: this.AWS_S3_BUCKET_NAME,
+          Key: imageName,
+          ContentType: 'image/png',
+          Body: imageBuffer,
+        })
+      );
+      return imageName;
+    } catch (ex) {
+      throw new Error('Could not save file at AWS S3 : ' + ex);
+    }
+  }
+
+  async getImage(fileName: string): Promise<PNGWithMetadata> {
+    if (!fileName) return null;
+    try {
+      if (!this.doesFileExistLocally(fileName)) {
+        const localFileStream = await this.saveFileFromCloud(fileName);
+        await new Promise<void>((resolve, reject) => {
+          localFileStream.on('finish', () => {
+            this.scheduleLocalFileDeletion(fileName);
+            this.checkLocalDiskUsageAndClean();
+            resolve();
+          });
+          localFileStream.on('error', (error) => {
+            this.logger.error('Error writing file:', error);
+            reject(error);
+          });
+        });
+      }
+      return PNG.sync.read(readFileSync(this.getImagePath(fileName)));
+    } catch (ex) {
+      this.logger.error(`Error from read : Cannot get image: ${fileName}. ${ex}`);
+    }
+  }
+
+  async saveFileFromCloud(fileName: string) {
+    const command = new GetObjectCommand({ Bucket: this.AWS_S3_BUCKET_NAME, Key: fileName });
+    const s3Response = await this.s3Client.send(command);
+    const fileStream = s3Response.Body as Readable;
+    const localFileStream = createWriteStream(this.getImagePath(fileName));
+    fileStream.pipe(localFileStream);
+    this.logger.log(`File feom AWS S3 saved at ${this.getImagePath(fileName)}`);
+    return localFileStream;
+  }
+
+  async deleteImage(imageName: string): Promise<boolean> {
+    if (!imageName) return false;
+    try {
+      await this.s3Client.send(new DeleteObjectCommand({ Bucket: this.AWS_S3_BUCKET_NAME, Key: imageName }));
+      return true;
+    } catch (error) {
+      this.logger.log(`Failed to delete file at AWS S3 for image ${imageName}:`, error);
+      return false; // Return `false` if an error occurs.
+    }
+  }
+}

src/shared/static/common-file-service.ts

@@ -0,0 +1,94 @@
+import path from 'path';
+import { unlink, mkdirSync, existsSync, unlinkSync, statSync, readdir } from 'fs';
+import { Logger } from '@nestjs/common/services/logger.service';
+
+export const IMAGE_PATH = 'imageUploads/';
+
+export class CommonFileService {
+  readonly DELETE_INTERVAL = 60 * 60 * 1000; // Check for deletions every hour
+  deletionQueue: { filePath: string; size: number }[] = [];
+  readonly MAX_DISK_USAGE: number = parseInt(process.env.MAX_TEMP_STORAGE_FOR_S3_DOWNLOAD) * 1 * 1024 * 1024;
+  logger: Logger;
+
+  generateNewImage(type: 'screenshot' | 'diff' | 'baseline'): { imageName: string; imagePath: string } {
+    const imageName = `${Date.now()}.${type}.png`;
+    return {
+      imageName,
+      imagePath: this.getImagePath(imageName),
+    };
+  }
+
+  getImagePath(imageName: string): string {
+    this.ensureDirectoryExistence(IMAGE_PATH);
+    return path.resolve(IMAGE_PATH, imageName);
+  }
+
+  private ensureDirectoryExistence(dir: string) {
+    const filePath = path.resolve(dir);
+    if (existsSync(filePath)) {
+      return true;
+    } else {
+      mkdirSync(dir, { recursive: true });
+      this.ensureDirectoryExistence(dir);
+    }
+  }
+
+  doesFileExistLocally(fileName: string) {
+    return existsSync(this.getImagePath(fileName));
+  }
+
+  scheduleLocalFileDeletion(fileName: string): void {
+    const filePath = this.getImagePath(fileName);
+    const fileSize = statSync(filePath).size;
+    this.deletionQueue.push({ filePath, size: fileSize });
+    this.logger.log(`Scheduled deletion for file: ${filePath} with size: ${fileSize} bytes`);
+  }
+
+  checkLocalDiskUsageAndClean(): void {
+    const totalDiskUsage = this.getTotalDiskUsage();
+    if (totalDiskUsage > this.MAX_DISK_USAGE) {
+      this.logger.log(`Disk usage exceeded. Triggering cleanup.`);
+      this.cleanupQueuedFiles();
+    }
+  }
+
+  getTotalDiskUsage(): number {
+    return this.deletionQueue.reduce((total, file) => total + file.size, 0);
+  }
+
+  cleanUpAllFiles() {
+    readdir(path.resolve(IMAGE_PATH), (err, files) => {
+      if (err) throw err;
+      for (const file of files) {
+        unlink(path.join(path.resolve(IMAGE_PATH), file), () => {});
+      }
+    });
+  }
+
+  cleanupQueuedFiles(): void {
+    const totalDiskUsage = this.getTotalDiskUsage();
+    this.logger.log(`Cleaning up files. Total disk usage: ${totalDiskUsage} bytes`);
+
+    // Delete files until the total disk usage is within the limit
+    let currentUsage = totalDiskUsage;
+
+    if (currentUsage > this.MAX_DISK_USAGE) {
+      // Sort files by the earliest (oldest) first for deletion
+      this.deletionQueue.sort((a, b) => statSync(a.filePath).birthtimeMs - statSync(b.filePath).birthtimeMs);
+
+      // Reduce the size by half
+      while (currentUsage > this.MAX_DISK_USAGE / 2 && this.deletionQueue.length > 0) {
+        const fileToDelete = this.deletionQueue.shift();
+        if (fileToDelete) {
+          try {
+            unlinkSync(fileToDelete.filePath);
+            currentUsage -= fileToDelete.size;
+            this.logger.log(`Deleted file: ${fileToDelete.filePath}`);
+          } catch (error) {
+            this.logger.log(`Failed to delete file: ${fileToDelete.filePath}. Error: ${error.message}`);
+          }
+        }
+      }
+    }
+  }
+}

src/shared/static/hard-disk.service.ts

@@ -0,0 +1,53 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { writeFileSync, readFileSync, unlink } from 'fs';
+import { PNG, PNGWithMetadata } from 'pngjs';
+
+import { StaticService } from './static-service.interface';
+import { CommonFileService } from './common-file-service';
+
+@Injectable()
+export class HardDiskService extends CommonFileService implements StaticService {
+  constructor() {
+    super();
+    this.logger = new Logger(HardDiskService.name);
+    this.logger.log('Local file system is used for file storage.');
+  }
+
+  saveFileFromCloud(fileName: string) {
+    throw new Error(`Download of file ${fileName} is not applicable for local files.`);
+  }
+
+  async saveImage(type: 'screenshot' | 'diff' | 'baseline', imageBuffer: Buffer): Promise<string> {
+    try {
+      new PNG().parse(imageBuffer);
+    } catch (ex) {
+      throw new Error('Cannot parse image as PNG file: ' + ex);
+    }
+
+    const { imageName, imagePath } = this.generateNewImage(type);
+    writeFileSync(imagePath, imageBuffer);
+    return imageName;
+  }
+
+  async getImage(fileName: string): Promise<PNGWithMetadata> {
+    if (!fileName) return null;
+    try {
+      return PNG.sync.read(readFileSync(this.getImagePath(fileName)));
+    } catch (ex) {
+      this.logger.error(`Error from read : Cannot get image: ${fileName}. ${ex}`);
+    }
+  }
+
+  async deleteImage(imageName: string): Promise<boolean> {
+    if (!imageName) return false;
+    return new Promise((resolvePromise) => {
+      unlink(this.getImagePath(imageName), (err) => {
+        if (err) {
+          this.logger.error(err);
+          resolvePromise(false);
+        }
+        resolvePromise(true);
+      });
+    });
+  }
+}

src/shared/static/static-service.interface.ts

@@ -0,0 +1,16 @@
+export interface StaticService {
+  // Below are service specific functions
+  saveImage(type: 'screenshot' | 'diff' | 'baseline', imageBuffer: Buffer);
+  getImage(fileName: string);
+  deleteImage(imageName: string);
+
+  doesFileExistLocally(fileName: string);
+  saveFileFromCloud(fileName: string);
+  scheduleLocalFileDeletion(fileName: string);
+  checkLocalDiskUsageAndClean();
+
+  getImagePath(imageName: string);
+  generateNewImage(imageName: string);
+}
+
+export const STATIC_SERVICE = 'FILE_SERVICE';