feat: enhance classifyImage functionality to support optional hash computation and update related types

abbottdev-kroll · abbottdev-kroll · commit 4363c90db1f1 · 2025-08-29T16:50:51.000+01:00
diff --git a/__tests__/unit/448x448.jpg b/__tests__/unit/448x448.jpg
diff --git a/__tests__/unit/578x478.jpg b/__tests__/unit/578x478.jpg
diff --git a/__tests__/unit/hashing.test.ts b/__tests__/unit/hashing.test.ts
@@ -0,0 +1,59 @@
+import { describe, it,} from 'vitest';
+import { computeHashesFromStream, HashType, ImageFormat, RequestEncoding } from '../../src';
+import fs from 'fs';
+
+describe('hashing', () => {
+  it('should compute MD5 and SHA1 hashes for a valid image', async ({expect}) => {
+    // Test implementation here
+    // Test implementation here
+    const imagePath = __dirname + '/448x448.jpg';
+
+    const imageStream = fs.createReadStream(imagePath);
+
+    //assert error is thrown for invalid image
+    const { data, format, md5, sha1 } = await computeHashesFromStream(imageStream,
+      RequestEncoding.UNCOMPRESSED,
+      ImageFormat.JPEG,
+      false,
+      [HashType.MD5, HashType.SHA1]
+    );
+
+    expect(data).toBeDefined();
+    expect(format).toBe(ImageFormat.JPEG);
+
+    expect(md5).toEqual('eb3a95fdd86ce28d9a63a68328783874');
+    expect(sha1).toEqual('b972b222bc91c457d904ebff16134dc79b67d1c9');
+  });
+
+  it('should compute MD5 and SHA1 hashes for a resized invalid image', async ({expect}) => {
+    // Test implementation here
+    // Test implementation here
+    const imagePath = __dirname + '/578x478.jpg';
+
+    const imageStream = fs.createReadStream(imagePath);
+
+    //assert error is thrown for invalid image
+    const { data, format, md5, sha1 } = await computeHashesFromStream(imageStream,
+      RequestEncoding.UNCOMPRESSED,
+      ImageFormat.JPEG,
+      true,
+      [HashType.MD5, HashType.SHA1]
+    );
+
+    expect(data).toBeDefined();
+    expect(format).toBe(ImageFormat.RAW_UINT8);
+
+    expect(md5).toEqual('d390b6cd7436fd41d1bcd005e7e3e652');
+    expect(sha1).toEqual('7f979ca35cfe390bd92f5dfa4a05919da76f0e43');
+  });
+
+  it('should throw an error for an invalid image', async ({expect}) => {
+    // Test implementation here
+    const imagePath = __dirname + '/578x478.jpg';
+
+    const imageStream = fs.createReadStream(imagePath);
+
+    //assert error is thrown for invalid image
+    await expect(computeHashesFromStream(imageStream)).rejects.toThrow('Image must be 448x448 pixels');
+  });
+});
diff --git a/__tests__/unit/main.test.ts b/__tests__/unit/main.test.ts
@@ -3,7 +3,8 @@ import { ClassificationOutput, ClassifierSdk, ClassifyImageInput, ImageFormat }
 import fs from 'fs';
 import { randomUUID } from 'crypto';
 
-describe('classifierHelper', () => {
+describe('ClassifierSdk smoke tests', () => {
+
   it('should listDeployments and return responses (smoke test)', async ({expect}) => {
       // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
       // You may want to mock the gRPC client for true unit testing.
@@ -34,7 +35,7 @@ describe('classifyImage function', () => {
   it('should classify 10 images in a single request and return responses (integration smoke test)', async ({expect, annotate}) => {
     // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
     // You may want to mock the gRPC client for true unit testing.
-    const imagePath = __dirname + '/Steamboat-willie.jpg';
+    const imagePath = __dirname + '/448x448.jpg';
     const sdk = new ClassifierSdk({
         deploymentId: process.env.VITE_ATHENA_DEPLOYMENT_ID,
         affiliate: process.env.VITE_ATHENA_AFFILIATE,
@@ -111,22 +112,21 @@ describe('classifyImage function', () => {
     const expectedOutputs = correlationIds.map(id =>(
       {
         correlationId: id,
-        classifications: expect.arrayContaining([
+        classifications: expect.toBeOneOf([expect.arrayContaining([
           {
             label: expect.any(String),
             weight: expect.any(Number)
           }
-        ])
+        ]), []])
       }
     ));
 
     expect(outputs).toMatchObject(expectedOutputs);
   }, 120000);
 
-  it('should classify Steamboat-willie.jpg with raw uint8 resize return responses (integration smoke test)', async ({expect, annotate}) => {
-    // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
-    // You may want to mock the gRPC client for true unit testing.
-    const imagePath = __dirname + '/Steamboat-willie.jpg';
+  it('should classify with raw uint8 resize return responses (integration smoke test)', async ({expect, annotate}) => {
+
+    const imagePath = __dirname + '/448x448.jpg';
     const sdk = new ClassifierSdk({
         deploymentId: process.env.VITE_ATHENA_DEPLOYMENT_ID,
         affiliate: process.env.VITE_ATHENA_AFFILIATE,
@@ -187,23 +187,23 @@ describe('classifyImage function', () => {
     expect(first).toMatchObject([
       {
         correlationId,
-        classifications: expect.arrayContaining([
+        classifications: expect.toBeOneOf([expect.arrayContaining([
           {
             label: expect.any(String),
             weight: expect.any(Number)
           }
-        ])
+        ]), []])
       } as ClassificationOutput
     ]);
 
     // Accept either a successful call or a connection error (for CI/dev convenience)
     expect(error).toBeUndefined();
   }, 120000);
 
-  it('should classify Steamboat-willie.jpg and return responses (integration smoke test)', async ({expect, annotate}) => {
+  it('should classify return responses (integration smoke test)', async ({expect, annotate}) => {
     // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
     // You may want to mock the gRPC client for true unit testing.
-    const imagePath = __dirname + '/Steamboat-willie.jpg';
+    const imagePath = __dirname + '/448x448.jpg';
     const sdk = new ClassifierSdk({
         deploymentId: process.env.VITE_ATHENA_DEPLOYMENT_ID,
         affiliate: process.env.VITE_ATHENA_AFFILIATE,
@@ -248,7 +248,7 @@ describe('classifyImage function', () => {
     const options: ClassifyImageInput = {
       imageStream,
       correlationId,
-      format: ImageFormat.JPEG
+      format: ImageFormat.JPEG,
     };
     try {
       await sdk.sendClassifyRequest(options);
@@ -264,12 +264,12 @@ describe('classifyImage function', () => {
     expect(first).toMatchObject([
       {
         correlationId,
-        classifications: expect.arrayContaining([
+        classifications: expect.toBeOneOf([expect.arrayContaining([
           {
             label: expect.any(String),
             weight: expect.any(Number)
           }
-        ])
+        ]), []])
       } as ClassificationOutput
     ]);
 
diff --git a/src/hashing.ts b/src/hashing.ts
@@ -1,7 +1,7 @@
 import { Readable } from 'stream';
 import crypto from 'crypto';
 import sharp from 'sharp';
-import { ImageFormat, RequestEncoding } from '.';
+import { HashType, ImageFormat, RequestEncoding } from '.';
 import brotli from 'brotli';
 import { buffer } from 'stream/consumers';
 
@@ -10,39 +10,55 @@ import { buffer } from 'stream/consumers';
  * @param stream Node.js readable stream (e.g., fs.createReadStream)
  * @param encoding Encoding type for the image data (default is UNCOMPRESSED)
  * @param imageFormat Format of the input image (default is UNSPECIFIED)
- * @param resize Whether to resize the image to 448x448 pixels (default is false
- * @returns {Promise<{md5: string, sha1: string, data: Buffer}>} Object containing MD5 hash, SHA1 hash, and resized image buffer
+ * @param resize Whether to resize the image to 448x448 pixels (default is false)
+ * @param hashes Array of hash types to compute (default is [MD5, SHA1])
+ * @returns {Promise<{md5?: string, sha1?: string, data: Buffer}>} Object containing MD5 hash, SHA1 hash, and resized image buffer
  */
 export async function computeHashesFromStream(
   stream: Readable,
   encoding: RequestEncoding = RequestEncoding.UNCOMPRESSED,
   imageFormat: ImageFormat = ImageFormat.UNSPECIFIED,
   resize: boolean = false,
-): Promise<{ md5: string; sha1: string; data: Buffer; format: ImageFormat }> {
+  hashes: HashType[] = [HashType.MD5, HashType.SHA1],
+): Promise<{ md5?: string; sha1?: string; data: Buffer; format: ImageFormat }> {
   const md5 = crypto.createHash('md5');
   const sha1 = crypto.createHash('sha1');
 
   let data: Buffer<ArrayBufferLike>;
 
-  stream.pipe(md5);
-  stream.pipe(sha1);
+  if (hashes.includes(HashType.MD5)) {
+    {
+      stream.pipe(md5);
+    }
+  }
+
+  if (hashes.includes(HashType.SHA1)) {
+    {
+      stream.pipe(sha1);
+    }
+  }
 
   if (resize) {
-    const resizer = sharp().resize(448, 448).raw({ depth: 'uint' });
+    const resizer = sharp().resize(448, 448).raw({ depth: 'char' });
     stream.pipe(resizer);
     data = await resizer.toBuffer();
     imageFormat = ImageFormat.RAW_UINT8;
   } else {
     data = await buffer(stream);
+    // use sharp to validate the image dimensions
+    const metadata = await sharp(data).metadata();
+    if (metadata.width !== 448 || metadata.height !== 448) {
+      throw new Error('Image must be 448x448 pixels');
+    }
   }
 
   if (encoding === RequestEncoding.BROTLI) {
     data = Buffer.from(await brotli.compress(data));
   }
 
   return {
-    md5: md5.digest('hex'),
-    sha1: sha1.digest('hex'),
+    md5: hashes.includes(HashType.MD5) ? md5.digest('hex') : undefined,
+    sha1: hashes.includes(HashType.SHA1) ? sha1.digest('hex') : undefined,
     data,
     format: imageFormat,
   };
diff --git a/src/index.ts b/src/index.ts
@@ -31,12 +31,15 @@ import { computeHashesFromStream } from './hashing';
  * @property imageStream The image data as a readable stream.
  * @property encoding Optional encoding type for the image data.
  * @property format Optional image format.
+ * @property resize Optional flag to resize the image.
+ * @property hashes Array of hash types to compute.
  */
 export type ClassifyImageInput = {
   affiliate?: string;
   correlationId?: string;
   imageStream: Readable;
   encoding?: ClassifyRequest['inputs'][number]['encoding'];
+  includeHashes?: HashType[];
 } & (ResizeImageInput | RawImageInput);
 
 export type ResizeImageInput = {
@@ -246,6 +249,7 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
         affiliate = this.options.affiliate,
         correlationId = randomUUID(),
         imageStream,
+        includeHashes = [HashType.MD5, HashType.SHA1],
         encoding = RequestEncoding.UNCOMPRESSED,
       } = options;
 
@@ -260,11 +264,18 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
         encoding,
         inputFormat,
         'resize' in options,
+        includeHashes,
       );
-      const hashes: ImageHash[] = [
-        { value: md5, type: HashType.MD5 },
-        { value: sha1, type: HashType.SHA1 },
-      ];
+
+      const hashes: ImageHash[] = [];
+
+      if (md5 && md5.trim() != '') {
+        hashes.push({ value: md5, type: HashType.MD5 });
+      }
+
+      if (sha1 && sha1.trim() != '') {
+        hashes.push({ value: sha1, type: HashType.SHA1 });
+      }
 
       processedInputs.push({
         affiliate,
