From 8fff46f31070730cb72b47d8a23a080dc04c737a Mon Sep 17 00:00:00 2001
From: Tanner Barlow <tanner.barlow12@gmail.com>
Date: Thu, 29 Aug 2019 15:55:44 -0700
Subject: [PATCH] fix: Clean up downloaded artifact after rollback

---
 src/services/rollbackService.test.ts | 15 +++++++++++++--
 src/services/rollbackService.ts      |  4 ++++
 src/test/mockFactory.ts              |  5 +++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/services/rollbackService.test.ts b/src/services/rollbackService.test.ts
index e5e0ce1c..502e03d7 100644
--- a/src/services/rollbackService.test.ts
+++ b/src/services/rollbackService.test.ts
@@ -5,6 +5,7 @@ import { ArmDeployment } from "../models/armTemplates";
 import { DeploymentConfig } from "../models/serverless";
 import { MockFactory } from "../test/mockFactory";
 import { RollbackService } from "./rollbackService";
+import fs from "fs";
 
 jest.mock("./azureBlobStorageService");
 import { AzureBlobStorageService } from "./azureBlobStorageService";
@@ -31,9 +32,10 @@ describe("Rollback Service", () => {
   const containerName = "deployment-artifacts";
   const artifactName = MockFactory.createTestDeployment().name.replace(
     configConstants.naming.suffix.deployment, configConstants.naming.suffix.artifact) + ".zip";
-  const artifactPath = `.serverless${path.sep}${artifactName}`
+  const artifactPath = path.join(".serverless", artifactName);
   const armDeployment: ArmDeployment = { template, parameters };
   const deploymentString = "deployments";
+  let unlinkSpy: jest.SpyInstance;
 
   function createOptions(timestamp?: string): Serverless.Options {
     return {
@@ -64,10 +66,12 @@ describe("Rollback Service", () => {
     ResourceService.prototype.listDeployments = jest.fn(() => Promise.resolve(deploymentString))
     AzureBlobStorageService.prototype.generateBlobSasTokenUrl = jest.fn(() => sasURL) as any;
     FunctionAppService.prototype.get = jest.fn(() => appStub) as any;
+    unlinkSpy = jest.spyOn(fs, "unlinkSync");
   });
 
   afterEach(() => {
     mockFs.restore();
+    unlinkSpy.mockRestore();
     jest.resetAllMocks();
   });
 
@@ -94,6 +98,11 @@ describe("Rollback Service", () => {
   });
 
   it("should deploy blob package directly to function app", async () => {
+    const fsConfig = {};
+    fsConfig[artifactPath] = "contents";
+    // Mocking the existence of the downloaded artifact because the downloadBinary
+    // method won't write to the mock file system
+    mockFs(fsConfig);
     const service = createService();
     await service.rollback();
     expect(AzureBlobStorageService.prototype.initialize).toBeCalled();
@@ -107,7 +116,8 @@ describe("Rollback Service", () => {
     expect(FunctionAppService.prototype.uploadZippedArfifactToFunctionApp).toBeCalledWith(
       appStub,
       artifactPath
-    )
+    );
+    expect(unlinkSpy).toBeCalledWith(artifactPath);
   });
 
   it("should deploy function app with SAS URL", async () => {
@@ -133,5 +143,6 @@ describe("Rollback Service", () => {
     );
     expect(FunctionAppService.prototype.get).not.toBeCalled();
     expect(FunctionAppService.prototype.uploadZippedArfifactToFunctionApp).not.toBeCalled();
+    expect(unlinkSpy).not.toBeCalled();
   });
 });
diff --git a/src/services/rollbackService.ts b/src/services/rollbackService.ts
index a8576938..dd9c911a 100644
--- a/src/services/rollbackService.ts
+++ b/src/services/rollbackService.ts
@@ -8,6 +8,7 @@ import { BaseService } from "./baseService";
 import { FunctionAppService } from "./functionAppService";
 import { ResourceService } from "./resourceService";
 import { ArmDeployment } from "../models/armTemplates";
+import fs from "fs";
 
 /**
  * Services for the Rollback Plugin
@@ -98,6 +99,9 @@ export class RollbackService extends BaseService {
     const functionAppService = new FunctionAppService(this.serverless, this.options);
     const functionApp = await functionAppService.get();
     await functionAppService.uploadZippedArfifactToFunctionApp(functionApp, artifactPath);
+    if (fs.existsSync(artifactPath)) {
+      fs.unlinkSync(artifactPath);
+    }
   }
 
   /**
diff --git a/src/test/mockFactory.ts b/src/test/mockFactory.ts
index d6c1a68f..b91e4f8f 100644
--- a/src/test/mockFactory.ts
+++ b/src/test/mockFactory.ts
@@ -16,6 +16,7 @@ import { ArmDeployment, ArmResourceTemplate, ArmTemplateProvisioningState } from
 import { ServicePrincipalEnvVariables } from "../models/azureProvider";
 import { Logger } from "../models/generic";
 import { ServerlessAzureConfig, ServerlessAzureProvider, ServerlessAzureFunctionConfig, ServerlessCliCommand } from "../models/serverless";
+import configConstants from "../config";
 
 function getAttribute(object: any, prop: string, defaultValue: any): any {
   if (object && object[prop]) {
@@ -174,7 +175,7 @@ export class MockFactory {
     const result = [];
     const originalTimestamp = +MockFactory.createTestTimestamp();
     for (let i = 0; i < count; i++) {
-      const name = (includeTimestamp) ? `deployment${i + 1}-t${originalTimestamp + i}` : `deployment${i + 1}`;
+      const name = (includeTimestamp) ? `${configConstants.naming.suffix.deployment}${i + 1}-t${originalTimestamp + i}` : `deployment${i + 1}`;
       result.push(
         MockFactory.createTestDeployment(name, i)
       )
@@ -194,7 +195,7 @@ export class MockFactory {
 
   public static createTestDeployment(name?: string, second: number = 0): DeploymentExtended {
     return {
-      name: name || `deployment1-t${MockFactory.createTestTimestamp()}`,
+      name: name || `${configConstants.naming.suffix.deployment}1-t${MockFactory.createTestTimestamp()}`,
       properties: {
         timestamp: new Date(2019, 1, 1, 0, 0, second),
         parameters: MockFactory.createTestParameters(),