handler.test.js

const { virusScan } = require("./handler");

var mockPutObjectTagging;
var mockPromise;

jest.mock("aws-sdk", () => {
  mockPromise = jest.fn();
  mockPutObjectTagging = jest.fn().mockReturnThis();

  return {
    S3: jest.fn().mockReturnValue({
      getObject: jest.fn().mockReturnThis(),
      putObjectTagging: mockPutObjectTagging,
      promise: mockPromise,
    }),
  };
});

var mockReadFileSync;
var mockWriteFileSync;
var mockUnlinkSync;

jest.mock("fs", () => {
  mockReadFileSync = jest.fn();
  mockWriteFileSync = jest.fn();
  mockUnlinkSync = jest.fn();

  return {
    readFileSync: mockReadFileSync,
    writeFileSync: mockWriteFileSync,
    unlinkSync: mockUnlinkSync,
  };
});

var mockExecSync;

jest.mock("child_process", () => {
  mockExecSync = jest.fn();

  return {
    execSync: mockExecSync,
  };
});

describe("handler.virusScan", () => {
  let mockValidEvent;

  beforeAll(() => {
    // console.log = jest.fn();

    mockValidEvent = {
      Records: [
        {
          s3: {
            bucket: { name: "bucket-name" },
            object: { key: "file-to-scan.jpeg" },
          },
        },
      ],
    };
  });

  it("calls an event with no event.Records", async () => {
    const event = {};
    await virusScan(event);

    // expect(console.log).toHaveBeenCalled();
    // expect(console.log).toHaveBeenCalledWith('Not an S3 event invocation!');
    expect(mockWriteFileSync).not.toHaveBeenCalled();
  });

  it("calls an event with no event.Records[].s3", async () => {
    const event = {
      Records: [
        {
          notS3: { foo: "bar" },
        },
      ],
    };

    await virusScan(event);

    // expect(console.log).toHaveBeenCalled();
    // expect(console.log).toHaveBeenCalledWith('Not an S3 Record!');
    expect(mockWriteFileSync).not.toHaveBeenCalled();
  });

  it("calls with a clean file", async () => {
    mockPromise.mockImplementationOnce(() => ({ Body: "body" }));

    await virusScan(mockValidEvent);

    expect(mockExecSync).toHaveBeenCalledWith(
      "clamscan /tmp/file-to-scan.jpeg"
    );
    expect(mockWriteFileSync).toHaveBeenCalledWith(
      "/tmp/file-to-scan.jpeg",
      "body"
    );
    expect(mockUnlinkSync).toHaveBeenCalledWith("/tmp/file-to-scan.jpeg");
    expect(mockPutObjectTagging).toHaveBeenCalledWith({
      Bucket: "bucket-name",
      Key: "file-to-scan.jpeg",
      Tagging: {
        TagSet: [{ Key: "av-status", Value: "clean" }],
      },
    });
  });

  it.only("calls with a virus file", async () => {
    mockPromise.mockImplementationOnce(() => ({ Body: "body" }));
    mockExecSync.mockImplementationOnce(() => {
      throw { code: 1 };
    });

    await virusScan(mockValidEvent);

    expect(mockExecSync).toHaveBeenCalledWith(
      "clamscan /tmp/file-to-scan.jpeg"
    );
    expect(mockWriteFileSync).toHaveBeenCalledWith(
      "/tmp/file-to-scan.jpeg",
      "body"
    );
    expect(mockUnlinkSync).toHaveBeenCalledWith("/tmp/file-to-scan.jpeg");
    expect(mockPutObjectTagging).toHaveBeenCalledWith({
      Bucket: "bucket-name",
      Key: "file-to-scan.jpeg",
      Tagging: {
        TagSet: [{ Key: "av-status", Value: "dirty" }],
      },
    });
  });
});