move tests from extension to listener test flush

Author: duncanista

src/metrics/extension.spec.ts

@@ -1,36 +1,24 @@
 import nock from "nock";
 
-import { isAgentRunning, flushExtension, AGENT_URL } from "./extension";
+import { isExtensionRunning, EXTENSION_URL } from "./extension";
 import mock from "mock-fs";
 
-describe("isAgentRunning", () => {
+describe("isExtensionRunning", () => {
   afterEach(() => {
     mock.restore();
   });
   it("returns true when agent exists and responds", async () => {
     mock({
       "/opt/extensions/datadog-agent": Buffer.from([0]),
     });
-    const ran = await isAgentRunning();
+    const ran = await isExtensionRunning();
     expect(ran).toBeTruthy();
   });
   it("returns false when agent doesn't exist", async () => {
     mock({});
-    const scope = nock(AGENT_URL).get("/lambda/hello").replyWithError("Unreachable");
-    const ran = await isAgentRunning();
+    const scope = nock(EXTENSION_URL).get("/lambda/hello").replyWithError("Unreachable");
+    const ran = await isExtensionRunning();
     expect(scope.isDone()).toBeFalsy();
     expect(ran).toBeFalsy();
   });
 });
-describe("flushExtension", () => {
-  it("calls flush on the agent", async () => {
-    const scope = nock(AGENT_URL).post("/lambda/flush", JSON.stringify({})).reply(200);
-    await flushExtension();
-    expect(scope.isDone()).toBeTruthy();
-  });
-  it("catches error when flush doesn't respond", async () => {
-    const scope = nock(AGENT_URL).post("/lambda/flush", JSON.stringify({})).replyWithError("Unavailable");
-    await flushExtension();
-    expect(scope.isDone()).toBeTruthy();
-  });
-});

src/metrics/listener.spec.ts

@@ -154,10 +154,50 @@ describe("MetricsListener", () => {
     });
     // jest.useFakeTimers();
 
-    listener.onStartInvocation({});
+    await listener.onStartInvocation({});
     listener.sendDistributionMetricWithDate("my-metric", 10, new Date(1584983836 * 1000), false, "tag:a", "tag:b");
     await listener.onCompleteInvocation();
 
     expect(spy).toHaveBeenCalledWith(`{"e":1584983836,"m":"my-metric","t":["tag:a","tag:b"],"v":10}\n`);
   });
+
+  describe("_localFlush", () => {
+    let listener: MetricsListener | undefined;
+    beforeEach(async () => {
+      mock({
+        "/opt/extensions/datadog-agent": Buffer.from([0]),
+      });
+      const distributionMock = jest.fn();
+      (StatsDClient as any).mockImplementation(() => {
+        return {
+          distribution: distributionMock,
+          close: (callback: any) => callback(undefined),
+        };
+      });
+      const kms = new MockKMS("kms-api-key-decrypted");
+      listener = new MetricsListener(kms as any, {
+        apiKey: "api-key",
+        apiKeyKMS: "",
+        enhancedMetrics: false,
+        logForwarding: true,
+        shouldRetryMetrics: false,
+        localTesting: true,
+        siteURL,
+      });
+      await listener.onStartInvocation({});
+    })
+    afterEach(() => {
+      listener = undefined;
+    })
+    it("calls flush on the agent", async () => {
+      const scope = nock(EXTENSION_URL).post("/lambda/flush", JSON.stringify({})).reply(200);
+      await listener!["_localFlush"]();
+      expect(scope.isDone()).toBeTruthy();
+    });
+    it("catches error when flush doesn't respond", async () => {
+      const scope = nock(EXTENSION_URL).post("/lambda/flush", JSON.stringify({})).replyWithError("Unavailable");
+      await listener!["_localFlush"]();
+      expect(scope.isDone()).toBeTruthy();
+    });
+  });
 });