Merge pull request #335 from Nick-1234531/DOT-5046/add-screenshot-sync-status

parthlambdatest · web-flow · commit 24872ac48c2f · 2025-08-27T17:11:22.000+05:30
Dot 5046/add screenshot sync status
diff --git a/src/lib/httpClient.ts b/src/lib/httpClient.ts
@@ -321,7 +321,7 @@ export default class httpClient {
         }, ctx.log)
     }
 
-    processSnapshot(ctx: Context, snapshot: ProcessedSnapshot, snapshotUuid: string,  discoveryErrors: DiscoveryErrors) {
+    processSnapshot(ctx: Context, snapshot: ProcessedSnapshot, snapshotUuid: string,  discoveryErrors: DiscoveryErrors, variantCount: number, sync: boolean = false) {
         return this.request({
             url: `/build/${ctx.build.id}/snapshot`,
             method: 'POST',
@@ -330,12 +330,14 @@ export default class httpClient {
                 name: snapshot.name,
                 url: snapshot.url,
                 snapshotUuid: snapshotUuid,
+                variantCount: variantCount,
                 test: {
                     type: ctx.testType,
                     source: 'cli'
                 },
                 doRemoteDiscovery: snapshot.options.doRemoteDiscovery,
                 discoveryErrors: discoveryErrors,
+                sync: sync
             }
         }, ctx.log)
     }
@@ -623,4 +625,14 @@ export default class httpClient {
             data: requestData
         }, ctx.log)
     }
+
+    getSnapshotStatus(snapshotName: string, snapshotUuid: string, ctx: Context): Promise<Record<string, any>> {
+        return this.request({
+            url: `/snapshot/status?buildId=${ctx.build.id}&snapshotName=${snapshotName}&snapshotUUID=${snapshotUuid}`,
+            method: 'GET',
+            headers: {
+                'Content-Type': 'application/json',
+            }
+        }, ctx.log);
+    }
 }
diff --git a/src/lib/schemaValidation.ts b/src/lib/schemaValidation.ts
@@ -501,6 +501,18 @@ const SnapshotSchema: JSONSchemaType<Snapshot> = {
                 sessionId: {
                     type: "string",
                     errorMessage: "Invalid snapshot options; sessionId must be a string"
+                },
+                contextId: {
+                    type: "string",
+                    errorMessage: "Invalid snapshot options; contextId must be a string"
+                },
+                sync: {
+                    type: "boolean",
+                    errorMessage: "Invalid snapshot options; sync must be a boolean"
+                },
+                timeout: {
+                    type: "number",
+                    errorMessage: "Invalid snapshot options; timeout must be a number"
                 }
             },
             additionalProperties: false
diff --git a/src/lib/server.ts b/src/lib/server.ts
@@ -41,6 +41,7 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
 			// Fetch sessionId from snapshot options if present
 			const sessionId = snapshot?.options?.sessionId;
 			let capsBuildId = ''
+			const contextId = snapshot?.options?.contextId;
 
 			if (sessionId) {
 				// Check if sessionId exists in the map
@@ -71,7 +72,23 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
 			}
 
 			ctx.testType = testType;
-			ctx.snapshotQueue?.enqueue(snapshot);
+			
+			if (contextId && !ctx.contextToSnapshotMap) {
+				ctx.contextToSnapshotMap = new Map();
+				ctx.log.debug(`Initialized empty context mapping map for contextId: ${contextId}`);
+			}
+			
+			if (contextId && ctx.contextToSnapshotMap) {
+				ctx.contextToSnapshotMap.set(contextId, 0);
+				ctx.log.debug(`Marking contextId as captured and added to queue: ${contextId}`);
+			}
+
+			if(contextId){
+				ctx.snapshotQueue?.enqueueFront(snapshot);
+			}else{
+				ctx.snapshotQueue?.enqueue(snapshot);	
+			}
+			
 			ctx.isSnapshotCaptured = true;
 			replyCode = 200;
 			replyBody = { data: { message: "success", warnings: [] }};
@@ -130,7 +147,113 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
 		reply.code(200).send({ status: 'Server is running', version: ctx.cliVersion });
 	});
 
-	
+	// Get snapshot status
+	server.get('/snapshot/status', opts, async (request, reply) => {
+		let replyCode: number;
+		let replyBody: Record<string, any>;
+
+
+		try {
+			ctx.log.debug(`request.query : ${JSON.stringify(request.query)}`);
+			const { contextId, pollTimeout, snapshotName } = request.query as { contextId: string, pollTimeout: number, snapshotName: string };
+			if (!contextId || !snapshotName) {
+				throw new Error('contextId and snapshotName are required parameters');
+			}
+
+			const timeoutDuration = pollTimeout*1000 || 30000; 
+
+			// Check if we have stored snapshot status for this contextId
+			if (ctx.contextToSnapshotMap?.has(contextId)) {
+				let contextStatus = ctx.contextToSnapshotMap.get(contextId);
+				
+				while (contextStatus==0) {
+					// Wait 5 seconds before next check
+					await new Promise(resolve => setTimeout(resolve, 5000));
+					
+					contextStatus = ctx.contextToSnapshotMap.get(contextId);
+				}
+
+				if(contextStatus==2){
+					throw new Error("Snapshot Failed");
+				}
+				
+				ctx.log.debug("Snapshot uploaded successfully");
+
+				// Poll external API until it returns 200 or timeout is reached
+				let lastExternalResponse: any = null; 
+				const startTime = Date.now(); 
+
+				while (true) {
+					try {
+						const externalResponse = await ctx.client.getSnapshotStatus(
+							snapshotName,
+							contextId,
+							ctx
+						);
+						
+						lastExternalResponse = externalResponse;
+
+						if (externalResponse.statusCode === 200) {
+							replyCode = 200;
+							replyBody = externalResponse.data;
+							return reply.code(replyCode).send(replyBody);
+						} else if (externalResponse.statusCode === 202 ) {
+							replyBody= externalResponse.data;
+							ctx.log.debug(`External API attempt: Still processing, Pending Screenshots ${externalResponse.snapshotCount}`);
+							await new Promise(resolve => setTimeout(resolve, 5000));
+						}else if(externalResponse.statusCode===404){
+							ctx.log.debug(`Snapshot still processing, not uploaded`);
+							await new Promise(resolve => setTimeout(resolve, 5000));
+						}else {
+							ctx.log.debug(`Unexpected response from external API: ${JSON.stringify(externalResponse)}`);
+							replyCode = 500;
+							replyBody = { 
+								error: { 
+									message: `Unexpected response from external API: ${externalResponse.statusCode}`,
+									externalApiStatus: externalResponse.statusCode
+								}
+							};
+							return reply.code(replyCode).send(replyBody);
+						}
+
+						ctx.log.debug(`timeoutDuration: ${timeoutDuration}`);
+						ctx.log.debug(`Time passed: ${Date.now() - startTime}`);
+
+						if (Date.now() - startTime > timeoutDuration) {
+							replyCode = 202; 
+							replyBody = {
+								data: {
+									message: 'Request timed out-> Snapshot still processing'
+								} 
+							};
+							return reply.code(replyCode).send(replyBody);
+						}
+
+					} catch (externalApiError: any) {
+						ctx.log.debug(`External API call failed: ${externalApiError.message}`);
+						replyCode = 500;
+						replyBody = { 
+							error: { 
+								message: `External API call failed: ${externalApiError.message}`
+							}
+						};
+						return reply.code(replyCode).send(replyBody);
+					}
+				}
+			} else {
+				// No snapshot found for this contextId
+				replyCode = 404;
+				replyBody = { error: { message: `No snapshot found for contextId: ${contextId}` } };
+				return reply.code(replyCode).send(replyBody);
+			}
+		} catch (error: any) {
+			ctx.log.debug(`snapshot status failed; ${error}`);
+			replyCode = 500;
+			replyBody = { error: { message: error.message } };
+			return reply.code(replyCode).send(replyBody);
+		}
+	});
+
 
 	await server.listen({ port: ctx.options.port });
 	// store server's address for SDK
diff --git a/src/lib/snapshotQueue.ts b/src/lib/snapshotQueue.ts
@@ -3,7 +3,7 @@ import { Snapshot, Context } from "../types.js";
 import constants from "./constants.js";
 import processSnapshot, {prepareSnapshot} from "./processSnapshot.js"
 import { v4 as uuidv4 } from 'uuid';
-import { startPolling, stopTunnelHelper } from "./utils.js";
+import { startPolling, stopTunnelHelper, calculateVariantCountFromSnapshot } from "./utils.js";
 
 const uploadDomToS3ViaEnv = process.env.USE_LAMBDA_INTERNAL || false;
 export default class Queue {
@@ -29,6 +29,16 @@ export default class Queue {
         }
     }
 
+    enqueueFront(item: Snapshot): void {
+        this.snapshots.unshift(item);
+        if (!this.ctx.config.delayedUpload) {
+            if (!this.processing) {
+                this.processing = true;
+                this.processNext();
+            }
+        }
+    }
+
     startProcessingfunc(): void {
         if (!this.processing) {
             this.processing = true;
@@ -129,6 +139,8 @@ export default class Queue {
         return drop;
     }
 
+    
+
     private filterVariants(snapshot: Snapshot, config: any): boolean {
         let allVariantsDropped = true;
 
@@ -273,6 +285,7 @@ export default class Queue {
                 this.processingSnapshot = snapshot?.name;
                 let drop = false;
 
+
                 if (this.ctx.isStartExec && !this.ctx.config.tunnel) {
                     this.ctx.log.info(`Processing Snapshot: ${snapshot?.name}`);
                 }
@@ -332,6 +345,7 @@ export default class Queue {
 
 
                     if (useCapsBuildId) {
+                        this.ctx.log.info(`Using cached buildId: ${capsBuildId}`);
                         if (useKafkaFlowCaps) {
                             const snapshotUuid = uuidv4();
                             let uploadDomToS3 = this.ctx.config.useLambdaInternal || uploadDomToS3ViaEnv;
@@ -378,18 +392,23 @@ export default class Queue {
                             }
                         }
                         if (this.ctx.build && this.ctx.build.useKafkaFlow) {
-                            const snapshotUuid = uuidv4();
+                            let snapshotUuid = uuidv4();
                             let snapshotUploadResponse
-                            let uploadDomToS3 = this.ctx.config.useLambdaInternal || uploadDomToS3ViaEnv;
-                            if (!uploadDomToS3) {
+                            if (snapshot?.options?.contextId && this.ctx.contextToSnapshotMap?.has(snapshot.options.contextId)) {
+                                snapshotUuid = snapshot.options.contextId;
+                             }
+                             let uploadDomToS3 = this.ctx.config.useLambdaInternal || uploadDomToS3ViaEnv;
+                             if (!uploadDomToS3) {
                                 this.ctx.log.debug(`Uploading dom to S3 for snapshot using presigned URL`);
                                 const presignedResponse = await this.ctx.client.getS3PresignedURLForSnapshotUpload(this.ctx, processedSnapshot.name, snapshotUuid);
                                 const uploadUrl = presignedResponse.data.url;
                                 snapshotUploadResponse = await this.ctx.client.uploadSnapshotToS3(this.ctx, uploadUrl, processedSnapshot);
-                            } else {
+                             } else {
                                 this.ctx.log.debug(`Uploading dom to S3 for snapshot using LSRS`);
                                 snapshotUploadResponse = await this.ctx.client.sendDomToLSRS(this.ctx, processedSnapshot, snapshotUuid);
-                            }
+                             }
+
+
                             if (!snapshotUploadResponse || Object.keys(snapshotUploadResponse).length === 0) {
                                 this.ctx.log.debug(`snapshot failed; Unable to upload dom to S3`);
                                 this.processedSnapshots.push({ name: snapshot?.name, error: `snapshot failed; Unable to upload dom to S3` });
@@ -403,11 +422,19 @@ export default class Queue {
                                         this.ctx.log.debug(`Closed browser context for snapshot ${snapshot.name}`);
                                     }
                                 }
+                                if(snapshot?.options?.contextId){
+                                    this.ctx.contextToSnapshotMap?.set(snapshot?.options?.contextId,2);
+                                }
                                 this.processNext();
                             } else {
-                                await this.ctx.client.processSnapshot(this.ctx, processedSnapshot, snapshotUuid, discoveryErrors);
+                                await this.ctx.client.processSnapshot(this.ctx, processedSnapshot, snapshotUuid, discoveryErrors,calculateVariantCountFromSnapshot(processedSnapshot, this.ctx.config),snapshot?.options?.sync);
+                                if(snapshot?.options?.contextId && this.ctx.contextToSnapshotMap?.has(snapshot.options.contextId)){
+                                    this.ctx.contextToSnapshotMap.set(snapshot.options.contextId, 1);
+                                }
+                                this.ctx.log.debug(`ContextId: ${snapshot?.options?.contextId} status set to uploaded`);
                             }
                         } else {
+                            this.ctx.log.info(`Uploading snapshot to S3`);
                             await this.ctx.client.uploadSnapshot(this.ctx, processedSnapshot,  discoveryErrors);
                         }
                         this.ctx.totalSnapshots++;
@@ -418,6 +445,9 @@ export default class Queue {
             } catch (error: any) {
                 this.ctx.log.debug(`snapshot failed; ${error}`);
                 this.processedSnapshots.push({ name: snapshot?.name, error: error.message });
+                if (snapshot?.options?.contextId && this.ctx.contextToSnapshotMap) {
+                    this.ctx.contextToSnapshotMap.set(snapshot.options.contextId, 2);
+                }
             }
             // Close open browser contexts and pages
             if (this.ctx.browser) {
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
@@ -456,5 +456,57 @@ export async function stopTunnelHelper(ctx: Context) {
     ctx.log.debug('Tunnel is Stopped ? ' + status);
 } 
 
+/**
+ * Calculate the number of variants for a snapshot based on the configuration
+ * @param config - The configuration object containing web and mobile settings
+ * @returns The total number of variants that would be generated
+ */
+export function calculateVariantCount(config: any): number {
+    let variantCount = 0;
+
+    // Calculate web variants
+    if (config.web) {
+        const browsers = config.web.browsers || [];
+        const viewports = config.web.viewports || [];
+        variantCount += browsers.length * viewports.length;
+    }
+
+    // Calculate mobile variants
+    if (config.mobile) {
+        const devices = config.mobile.devices || [];
+        variantCount += devices.length;
+    }
+
+    return variantCount;
+}
+
+/**
+ * Calculate the number of variants for a snapshot based on snapshot-specific options
+ * @param snapshot - The snapshot object with options
+ * @param globalConfig - The global configuration object (fallback)
+ * @returns The total number of variants that would be generated
+ */
+export function calculateVariantCountFromSnapshot(snapshot: any, globalConfig?: any): number {
+    let variantCount = 0;
+    
+
+    // Check snapshot-specific web options
+    if (snapshot.options?.web) {
+        const browsers = snapshot.options.web.browsers || [];
+        const viewports = snapshot.options.web.viewports || [];
+        variantCount += browsers.length * viewports.length;
+    }
 
+    // Check snapshot-specific mobile options
+    if (snapshot.options?.mobile) {
+        const devices = snapshot.options.mobile.devices || [];
+        variantCount += devices.length;
+    }
+
+    // Fallback to global config if no snapshot-specific options
+    if (variantCount === 0 && globalConfig) {
+        variantCount = calculateVariantCount(globalConfig);
+    }
 
+    return variantCount;
+} 
diff --git a/src/types.ts b/src/types.ts
@@ -87,6 +87,7 @@ export interface Context {
     mergeBuildTargetId?: string;
     mergeByBranch?: boolean;
     mergeByBuild?: boolean;
+    contextToSnapshotMap?: Map<string, number>;
 }
 
 export interface Env {
@@ -147,6 +148,8 @@ export interface Snapshot {
         loadDomContent?: boolean;
         ignoreType?: string[],
         sessionId?: string
+        sync?: boolean;
+        contextId?: string;
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -87,6 +87,7 @@ export interface Context {`
`87`	`87`	`mergeBuildTargetId?: string;`
`88`	`88`	`mergeByBranch?: boolean;`
`89`	`89`	`mergeByBuild?: boolean;`
	`90`	`+ contextToSnapshotMap?: Map<string, number>;`
`90`	`91`	`}`
`91`	`92`
`92`	`93`	`export interface Env {`
`@@ -147,6 +148,8 @@ export interface Snapshot {`
`147`	`148`	`loadDomContent?: boolean;`
`148`	`149`	`ignoreType?: string[],`
`149`	`150`	`sessionId?: string`
	`151`	`+ sync?: boolean;`
	`152`	`+ contextId?: string;`
`150`	`153`	`}`
`151`	`154`	`}`
`152`	`155`