Fix unit tests

x-pack/plugins/fleet/server/services/agents/acks.test.ts

@@ -107,16 +107,21 @@ describe('test agent acks services', () => {
       ]
     );
     expect(mockSavedObjectsClient.bulkUpdate).not.toBeCalled();
-    expect(mockSavedObjectsClient.update).toBeCalled();
-    expect(mockSavedObjectsClient.update.mock.calls[0]).toMatchInlineSnapshot(`
+    expect(mockElasticsearchClient.update).toBeCalled();
+    expect(mockElasticsearchClient.update.mock.calls[0]).toMatchInlineSnapshot(`
       Array [
-        "fleet-agents",
-        "id",
         Object {
-          "packages": Array [
-            "system",
-          ],
-          "policy_revision": 4,
+          "body": Object {
+            "doc": Object {
+              "packages": Array [
+                "system",
+              ],
+              "policy_revision_idx": 4,
+            },
+          },
+          "id": "id",
+          "index": ".fleet-agents",
+          "refresh": "wait_for",
         },
       ]
     `);
@@ -169,16 +174,21 @@ describe('test agent acks services', () => {
       ]
     );
     expect(mockSavedObjectsClient.bulkUpdate).not.toBeCalled();
-    expect(mockSavedObjectsClient.update).toBeCalled();
-    expect(mockSavedObjectsClient.update.mock.calls[0]).toMatchInlineSnapshot(`
+    expect(mockElasticsearchClient.update).toBeCalled();
+    expect(mockElasticsearchClient.update.mock.calls[0]).toMatchInlineSnapshot(`
       Array [
-        "fleet-agents",
-        "id",
         Object {
-          "packages": Array [
-            "system",
-          ],
-          "policy_revision": 4,
+          "body": Object {
+            "doc": Object {
+              "packages": Array [
+                "system",
+              ],
+              "policy_revision_idx": 4,
+            },
+          },
+          "id": "id",
+          "index": ".fleet-agents",
+          "refresh": "wait_for",
         },
       ]
     `);

x-pack/plugins/fleet/server/services/agents/checkin/state_new_actions.test.ts

@@ -122,6 +122,7 @@ describe('test agent checkin new action services', () => {
       current_error_events: [],
       packages: [],
       enrolled_at: '2020-03-14T19:45:02.620Z',
+      default_api_key: 'MOCK_API_KEY',
     };
     const mockPolicyAction: AgentPolicyAction = {
       id: 'action1',

x-pack/plugins/fleet/server/services/agents/reassign.test.ts

@@ -7,101 +7,101 @@
 
 import { elasticsearchServiceMock, savedObjectsClientMock } from 'src/core/server/mocks';
 import type { SavedObject } from 'kibana/server';
-import type { Agent, AgentPolicy } from '../../types';
+import type { AgentPolicy } from '../../types';
 import { AgentReassignmentError } from '../../errors';
 import { reassignAgent, reassignAgents } from './reassign';
 
-const agentInManagedSO = {
-  id: 'agent-in-managed-policy',
-  attributes: { policy_id: 'managed-agent-policy' },
-} as SavedObject<Agent>;
-const agentInManagedSO2 = {
-  id: 'agent-in-managed-policy2',
-  attributes: { policy_id: 'managed-agent-policy' },
-} as SavedObject<Agent>;
-const agentInUnmanagedSO = {
-  id: 'agent-in-unmanaged-policy',
-  attributes: { policy_id: 'unmanaged-agent-policy' },
-} as SavedObject<Agent>;
-const agentInUnmanagedSO2 = {
-  id: 'agent-in-unmanaged-policy2',
-  attributes: { policy_id: 'unmanaged-agent-policy' },
-} as SavedObject<Agent>;
+const agentInManagedDoc = {
+  _id: 'agent-in-managed-policy',
+  _source: { policy_id: 'managed-agent-policy' },
+};
+const agentInManagedDoc2 = {
+  _id: 'agent-in-managed-policy2',
+  _source: { policy_id: 'managed-agent-policy' },
+};
+const agentInUnmanagedDoc = {
+  _id: 'agent-in-unmanaged-policy',
+  _source: { policy_id: 'unmanaged-agent-policy' },
+};
 const unmanagedAgentPolicySO = {
   id: 'unmanaged-agent-policy',
   attributes: { is_managed: false },
 } as SavedObject<AgentPolicy>;
+const unmanagedAgentPolicySO2 = {
+  id: 'unmanaged-agent-policy-2',
+  attributes: { is_managed: false },
+} as SavedObject<AgentPolicy>;
 const managedAgentPolicySO = {
   id: 'managed-agent-policy',
   attributes: { is_managed: true },
 } as SavedObject<AgentPolicy>;
 
 describe('reassignAgent (singular)', () => {
   it('can reassign from unmanaged policy to unmanaged', async () => {
-    const soClient = createClientMock();
-    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
-    await reassignAgent(soClient, esClient, agentInUnmanagedSO.id, agentInUnmanagedSO2.id);
+    const { soClient, esClient } = createClientsMock();
+    await reassignAgent(soClient, esClient, agentInUnmanagedDoc._id, unmanagedAgentPolicySO.id);
 
     // calls ES update with correct values
-    expect(soClient.update).toBeCalledTimes(1);
-    const calledWith = soClient.update.mock.calls[0];
-    expect(calledWith[1]).toBe(agentInUnmanagedSO.id);
-    expect(calledWith[2]).toHaveProperty('policy_id', agentInUnmanagedSO2.id);
+    expect(esClient.update).toBeCalledTimes(1);
+    const calledWith = esClient.update.mock.calls[0];
+    expect(calledWith[0]?.id).toBe(agentInUnmanagedDoc._id);
+    // @ts-expect-error
+    expect(calledWith[0]?.body?.doc).toHaveProperty('policy_id', unmanagedAgentPolicySO.id);
   });
 
   it('cannot reassign from unmanaged policy to managed', async () => {
-    const soClient = createClientMock();
-    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+    const { soClient, esClient } = createClientsMock();
     await expect(
-      reassignAgent(
-        soClient,
-        esClient,
-        agentInUnmanagedSO.id,
-        agentInManagedSO.attributes.policy_id!
-      )
+      reassignAgent(soClient, esClient, agentInUnmanagedDoc._id, managedAgentPolicySO.id)
     ).rejects.toThrowError(AgentReassignmentError);
 
     // does not call ES update
-    expect(soClient.update).toBeCalledTimes(0);
+    expect(esClient.update).toBeCalledTimes(0);
   });
 
   it('cannot reassign from managed policy', async () => {
-    const soClient = createClientMock();
-    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+    const { soClient, esClient } = createClientsMock();
     await expect(
-      reassignAgent(soClient, esClient, agentInManagedSO.id, agentInManagedSO2.id)
+      reassignAgent(soClient, esClient, agentInManagedDoc._id, unmanagedAgentPolicySO.id)
     ).rejects.toThrowError(AgentReassignmentError);
     // does not call ES update
-    expect(soClient.update).toBeCalledTimes(0);
+    expect(esClient.update).toBeCalledTimes(0);
 
     await expect(
-      reassignAgent(soClient, esClient, agentInManagedSO.id, agentInUnmanagedSO.id)
+      reassignAgent(soClient, esClient, agentInManagedDoc._id, managedAgentPolicySO.id)
     ).rejects.toThrowError(AgentReassignmentError);
     // does not call ES update
-    expect(soClient.update).toBeCalledTimes(0);
+    expect(esClient.update).toBeCalledTimes(0);
   });
 });
 
 describe('reassignAgents (plural)', () => {
   it('agents in managed policies are not updated', async () => {
-    const soClient = createClientMock();
-    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
-    const idsToReassign = [agentInUnmanagedSO.id, agentInManagedSO.id, agentInUnmanagedSO.id];
-    await reassignAgents(soClient, esClient, { agentIds: idsToReassign }, agentInUnmanagedSO.id);
+    const { soClient, esClient } = createClientsMock();
+    const idsToReassign = [agentInUnmanagedDoc._id, agentInManagedDoc._id, agentInManagedDoc2._id];
+    await reassignAgents(
+      soClient,
+      esClient,
+      { agentIds: idsToReassign },
+      unmanagedAgentPolicySO2.id
+    );
 
     // calls ES update with correct values
-    const calledWith = soClient.bulkUpdate.mock.calls[0][0];
-    const expectedResults = [agentInUnmanagedSO.id, agentInUnmanagedSO.id];
-    expect(calledWith.length).toBe(expectedResults.length); // only 2 are unmanaged
-    expect(calledWith.map(({ id }) => id)).toEqual(expectedResults);
+    const calledWith = esClient.bulk.mock.calls[0][0];
+    // only 1 are unmanaged and bulk write two line per update
+    // @ts-expect-error
+    expect(calledWith.body.length).toBe(2);
+    // @ts-expect-error
+    expect(calledWith.body[0].update._id).toEqual(agentInUnmanagedDoc._id);
   });
 });
 
-function createClientMock() {
+function createClientsMock() {
   const soClientMock = savedObjectsClientMock.create();
 
   // need to mock .create & bulkCreate due to (bulk)createAgentAction(s) in reassignAgent(s)
-  soClientMock.create.mockResolvedValue(agentInUnmanagedSO);
+  // @ts-expect-error
+  soClientMock.create.mockResolvedValue({ attributes: { agent_id: 'test' } });
   soClientMock.bulkCreate.mockImplementation(async ([{ type, attributes }]) => {
     return {
       saved_objects: [await soClientMock.create(type, attributes)],
@@ -110,26 +110,48 @@ function createClientMock() {
   soClientMock.bulkUpdate.mockResolvedValue({
     saved_objects: [],
   });
-
   soClientMock.get.mockImplementation(async (_, id) => {
     switch (id) {
       case unmanagedAgentPolicySO.id:
         return unmanagedAgentPolicySO;
       case managedAgentPolicySO.id:
         return managedAgentPolicySO;
-      case agentInManagedSO.id:
-        return agentInManagedSO;
-      case agentInUnmanagedSO.id:
+      case unmanagedAgentPolicySO2.id:
+        return unmanagedAgentPolicySO2;
       default:
-        return agentInUnmanagedSO;
+        throw new Error('Not found');
     }
   });
-
   soClientMock.bulkGet.mockImplementation(async (options) => {
     return {
       saved_objects: await Promise.all(options!.map(({ type, id }) => soClientMock.get(type, id))),
     };
   });
 
-  return soClientMock;
+  const esClientMock = elasticsearchServiceMock.createClusterClient().asInternalUser;
+  // @ts-expect-error
+  esClientMock.mget.mockImplementation(async () => {
+    return {
+      body: {
+        docs: [agentInManagedDoc, agentInUnmanagedDoc, agentInManagedDoc2],
+      },
+    };
+  });
+  // @ts-expect-error
+  esClientMock.get.mockImplementation(async ({ id }) => {
+    switch (id) {
+      case agentInManagedDoc._id:
+        return { body: agentInManagedDoc };
+      case agentInUnmanagedDoc._id:
+        return { body: agentInUnmanagedDoc };
+      default:
+        throw new Error('Not found');
+    }
+  });
+  // @ts-expect-error
+  esClientMock.bulk.mockResolvedValue({
+    body: { items: [] },
+  });
+
+  return { soClient: soClientMock, esClient: esClientMock };
 }