Add code and tests for updated force param for agent unenroll api

x-pack/plugins/fleet/server/routes/agent/unenroll_handler.ts

@@ -28,11 +28,10 @@ export const postAgentUnenrollHandler: RequestHandler<
   const soClient = context.core.savedObjects.client;
   const esClient = context.core.elasticsearch.client.asInternalUser;
   try {
-    if (request.body?.revoke === true) {
-      await AgentService.forceUnenrollAgent(soClient, esClient, request.params.agentId);
-    } else {
-      await AgentService.unenrollAgent(soClient, esClient, request.params.agentId);
-    }
+    await AgentService.unenrollAgent(soClient, esClient, request.params.agentId, {
+      force: request.body?.force,
+      revoke: request.body?.revoke,
+    });
 
     const body: PostAgentUnenrollResponse = {};
     return response.ok({ body });
@@ -63,6 +62,7 @@ export const postBulkAgentsUnenrollHandler: RequestHandler<
     const results = await AgentService.unenrollAgents(soClient, esClient, {
       ...agentOptions,
       revoke: request.body?.revoke,
+      force: request.body?.force,
     });
     const body = results.items.reduce<PostBulkAgentUnenrollResponse>((acc, so) => {
       acc[so.id] = {

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

@@ -46,7 +46,7 @@ describe('unenrollAgent (singular)', () => {
     expect(calledWith[0]?.body).toHaveProperty('doc.unenrollment_started_at');
   });
 
-  it('cannot unenroll from managed policy', async () => {
+  it('cannot unenroll from managed policy by default', async () => {
     const { soClient, esClient } = createClientMock();
     await expect(unenrollAgent(soClient, esClient, agentInManagedDoc._id)).rejects.toThrowError(
       AgentUnenrollmentError
@@ -54,6 +54,35 @@ describe('unenrollAgent (singular)', () => {
     // does not call ES update
     expect(esClient.update).toBeCalledTimes(0);
   });
+
+  it('cannot unenroll from managed policy with revoke=true', async () => {
+    const { soClient, esClient } = createClientMock();
+    await expect(
+      unenrollAgent(soClient, esClient, agentInManagedDoc._id, { revoke: true })
+    ).rejects.toThrowError(AgentUnenrollmentError);
+    // does not call ES update
+    expect(esClient.update).toBeCalledTimes(0);
+  });
+
+  it('can unenroll from managed policy with force=true', async () => {
+    const { soClient, esClient } = createClientMock();
+    await unenrollAgent(soClient, esClient, agentInManagedDoc._id, { force: true });
+    // calls ES update with correct values
+    expect(esClient.update).toBeCalledTimes(1);
+    const calledWith = esClient.update.mock.calls[0];
+    expect(calledWith[0]?.id).toBe(agentInManagedDoc._id);
+    expect(calledWith[0]?.body).toHaveProperty('doc.unenrollment_started_at');
+  });
+
+  it('can unenroll from managed policy with force=true and revoke=true', async () => {
+    const { soClient, esClient } = createClientMock();
+    await unenrollAgent(soClient, esClient, agentInManagedDoc._id, { force: true, revoke: true });
+    // calls ES update with correct values
+    expect(esClient.update).toBeCalledTimes(1);
+    const calledWith = esClient.update.mock.calls[0];
+    expect(calledWith[0]?.id).toBe(agentInManagedDoc._id);
+    expect(calledWith[0]?.body).toHaveProperty('doc.unenrolled_at');
+  });
 });
 
 describe('unenrollAgents (plural)', () => {
@@ -68,13 +97,12 @@ describe('unenrollAgents (plural)', () => {
       .filter((i: any) => i.update !== undefined)
       .map((i: any) => i.update._id);
     const docs = calledWith?.body.filter((i: any) => i.doc).map((i: any) => i.doc);
-    expect(ids).toHaveLength(2);
     expect(ids).toEqual(idsToUnenroll);
     for (const doc of docs) {
       expect(doc).toHaveProperty('unenrollment_started_at');
     }
   });
-  it('cannot unenroll from a managed policy', async () => {
+  it('cannot unenroll from a managed policy by default', async () => {
     const { soClient, esClient } = createClientMock();
 
     const idsToUnenroll = [
@@ -91,12 +119,116 @@ describe('unenrollAgents (plural)', () => {
       .filter((i: any) => i.update !== undefined)
       .map((i: any) => i.update._id);
     const docs = calledWith?.body.filter((i: any) => i.doc).map((i: any) => i.doc);
-    expect(ids).toHaveLength(onlyUnmanaged.length);
     expect(ids).toEqual(onlyUnmanaged);
     for (const doc of docs) {
       expect(doc).toHaveProperty('unenrollment_started_at');
     }
   });
+
+  it('cannot unenroll from a managed policy with revoke=true', async () => {
+    const { soClient, esClient } = createClientMock();
+
+    const idsToUnenroll = [
+      agentInUnmanagedDoc._id,
+      agentInManagedDoc._id,
+      agentInUnmanagedDoc2._id,
+    ];
+
+    const unenrolledResponse = await unenrollAgents(soClient, esClient, {
+      agentIds: idsToUnenroll,
+      revoke: true,
+    });
+
+    expect(unenrolledResponse.items).toMatchObject([
+      {
+        id: 'agent-in-unmanaged-policy',
+        success: true,
+      },
+      {
+        id: 'agent-in-managed-policy',
+        success: false,
+      },
+      {
+        id: 'agent-in-unmanaged-policy2',
+        success: true,
+      },
+    ]);
+
+    // calls ES update with correct values
+    const onlyUnmanaged = [agentInUnmanagedDoc._id, agentInUnmanagedDoc2._id];
+    const calledWith = esClient.bulk.mock.calls[0][0];
+    const ids = calledWith?.body
+      .filter((i: any) => i.update !== undefined)
+      .map((i: any) => i.update._id);
+    const docs = calledWith?.body.filter((i: any) => i.doc).map((i: any) => i.doc);
+    expect(ids).toEqual(onlyUnmanaged);
+    for (const doc of docs) {
+      expect(doc).toHaveProperty('unenrolled_at');
+    }
+  });
+
+  it('can unenroll from managed policy with force=true', async () => {
+    const { soClient, esClient } = createClientMock();
+    const idsToUnenroll = [
+      agentInUnmanagedDoc._id,
+      agentInManagedDoc._id,
+      agentInUnmanagedDoc2._id,
+    ];
+    await unenrollAgents(soClient, esClient, { agentIds: idsToUnenroll, force: true });
+
+    // calls ES update with correct values
+    const calledWith = esClient.bulk.mock.calls[1][0];
+    const ids = calledWith?.body
+      .filter((i: any) => i.update !== undefined)
+      .map((i: any) => i.update._id);
+    const docs = calledWith?.body.filter((i: any) => i.doc).map((i: any) => i.doc);
+    expect(ids).toEqual(idsToUnenroll);
+    for (const doc of docs) {
+      expect(doc).toHaveProperty('unenrollment_started_at');
+    }
+  });
+
+  it('can unenroll from managed policy with force=true and revoke=true', async () => {
+    const { soClient, esClient } = createClientMock();
+
+    const idsToUnenroll = [
+      agentInUnmanagedDoc._id,
+      agentInManagedDoc._id,
+      agentInUnmanagedDoc2._id,
+    ];
+
+    const unenrolledResponse = await unenrollAgents(soClient, esClient, {
+      agentIds: idsToUnenroll,
+      revoke: true,
+      force: true,
+    });
+
+    expect(unenrolledResponse.items).toMatchObject([
+      {
+        id: 'agent-in-unmanaged-policy',
+        success: true,
+      },
+      {
+        id: 'agent-in-managed-policy',
+        success: true,
+      },
+      {
+        id: 'agent-in-unmanaged-policy2',
+        success: true,
+      },
+    ]);
+
+    // calls ES update with correct values
+    const calledWith = esClient.bulk.mock.calls[0][0];
+    const ids = calledWith?.body
+      .filter((i: any) => i.update !== undefined)
+      .map((i: any) => i.update._id);
+    const docs = calledWith?.body.filter((i: any) => i.doc).map((i: any) => i.doc);
+    expect(ids).toEqual(idsToUnenroll);
+    for (const doc of docs) {
+      expect(doc).toHaveProperty('unenrolled_at');
+    }
+  });
 });
 
 function createClientMock() {

x-pack/plugins/fleet/server/services/agents/unenroll.ts

@@ -39,10 +39,18 @@ async function unenrollAgentIsAllowed(
 export async function unenrollAgent(
   soClient: SavedObjectsClientContract,
   esClient: ElasticsearchClient,
-  agentId: string
+  agentId: string,
+  options?: {
+    force?: boolean;
+    revoke?: boolean;
+  }
 ) {
-  await unenrollAgentIsAllowed(soClient, esClient, agentId);
-
+  if (!options?.force) {
+    await unenrollAgentIsAllowed(soClient, esClient, agentId);
+  }
+  if (options?.revoke) {
+    return forceUnenrollAgent(soClient, esClient, agentId);
+  }
   const now = new Date().toISOString();
   await createAgentAction(soClient, esClient, {
     agent_id: agentId,
@@ -57,7 +65,10 @@ export async function unenrollAgent(
 export async function unenrollAgents(
   soClient: SavedObjectsClientContract,
   esClient: ElasticsearchClient,
-  options: GetAgentsOptions & { revoke?: boolean }
+  options: GetAgentsOptions & {
+    force?: boolean;
+    revoke?: boolean;
+  }
 ): Promise<{ items: BulkActionResult[] }> {
   // start with all agents specified
   const givenAgents = await getAgents(esClient, options);
@@ -76,15 +87,17 @@ export async function unenrollAgents(
     )
   );
   const outgoingErrors: Record<Agent['id'], Error> = {};
-  const agentsToUpdate = agentResults.reduce<Agent[]>((agents, result, index) => {
-    if (result.status === 'fulfilled') {
-      agents.push(result.value);
-    } else {
-      const id = givenAgents[index].id;
-      outgoingErrors[id] = result.reason;
-    }
-    return agents;
-  }, []);
+  const agentsToUpdate = options.force
+    ? agentsEnrolled
+    : agentResults.reduce<Agent[]>((agents, result, index) => {
+        if (result.status === 'fulfilled') {
+          agents.push(result.value);
+        } else {
+          const id = givenAgents[index].id;
+          outgoingErrors[id] = result.reason;
+        }
+        return agents;
+      }, []);
 
   const now = new Date().toISOString();
   if (options.revoke) {

x-pack/plugins/fleet/server/types/rest_spec/agent.ts

@@ -172,7 +172,8 @@ export const PostAgentUnenrollRequestSchema = {
   }),
   body: schema.nullable(
     schema.object({
-      revoke: schema.boolean(),
+      force: schema.maybe(schema.boolean()),
+      revoke: schema.maybe(schema.boolean()),
     })
   ),
 };
@@ -180,6 +181,7 @@ export const PostAgentUnenrollRequestSchema = {
 export const PostBulkAgentUnenrollRequestSchema = {
   body: schema.object({
     agents: schema.oneOf([schema.arrayOf(schema.string()), schema.string()]),
+    force: schema.maybe(schema.boolean()),
     revoke: schema.maybe(schema.boolean()),
   }),
 };