[Fleet] Managed Agent Policy (#88688)

## Summary
Introduces the concept of a managed agent policy. Resolves most of the acceptance criteria from #76843. Remaining to be done in follow up PRs

- [x] Define hosted Agent Policy concept in Fleet.
    - [x] Flag in policy? **_yes, added `is_managed: boolean`_ in agent policy SO**
    - [x] Should not built only for cloud, an admin should be able to set theses restrictions.
    - [x] We should have an API to configure it _**Can `POST` and `PUT` to  `/api/fleet/agent_policies/{policy_id}`**_
    - [x] Integration should be editable, we expect integration author to do the right thing and limit what can be edited.
- [x] Research if we can ensure the right behavior of Hosted Agent policy and restrict the super user.
- [ ] Capabilities restrictions
  - [ ] An Agent enrolled in an Hosted Agent policy should not be able to be upgraded.
  - [x] An Agent enrolled in an Hosted Agent policy should not be able to be unenrolled.
  - [ ] No Agents cannot be enrolled into this policy by the user.
      - Hide the enrollment key?
      - Need to figure out the workflow.
  - [x] An Agent enrolled in an Hosted Agent policy should not be able to be reassigned to a different configuration.
- [x] As a user I should be prevented to do theses action. _**No user-level checks. Only Agent Policy. No UI changes, but API errors are shown for failed actions like reassigning**_
- [x] As an API user I should receive error messages.
- [x] If making a single "flag" is easier/faster let's do it.  _**Currently single `is_managed` property on agent policy SO.**_

Checks are implemented in service layer (is agent enrolled in a managed policy?)

No UI-specific changes added but UI is affected because HTTP requests (like `api/fleet/agents/{agentId}/reassign`) can fail. See screenshots below.

Tests at service (`yarn test:jest`) and http (`yarn test ftr`) layers for each of create policy, update policy, unenroll agent, and reassign agent

Bulk actions currently filter out restricted items. A follow-up PR will change them to throw an error and cause the request to fail.


## Managed Policy
Can create (`POST`) and update (`PUT`) an agent policy with an `is_managed` property. Each new saved object will have an `is_managed` property (default `false`)

<details><summary>HTTP commands</summary>

#### Create (`is_managed: false` by default)
```
 curl --user elastic:changeme -X POST localhost:5601/api/fleet/agent_policies -H 'Content-Type: application/json' -d'{ "name": "User created policy", "namespace": "default"}' -H 'kbn-xsrf: true'
{"item":{"id":"edc236a0-5cbb-11eb-ab2c-0134aecb4ce8","name":"User created policy","namespace":"default","is_managed":false,"revision":1,"updated_at":"2021-01-22T14:12:58.250Z","updated_by":"elastic"}}
```

#### Create with `is_managed: true`
```
 curl --user elastic:changeme -X POST localhost:5601/api/fleet/agent_policies -H 'Content-Type: application/json' -d'{ "name": "User created policy", "namespace": "default"}' -H 'kbn-xsrf: true'
{"item":{"id":"67c785b0-662e-11eb-bf6b-4790dc0178c0","name":"User created policy","namespace":"default","is_managed":false,"revision":1,"updated_at":"2021-02-03T14:45:06.059Z","updated_by":"elastic"}}
```

#### Update with `is_managed: true`
```
 curl --user elastic:changeme -X PUT  -H 'Content-Type: application/json' -H 'kbn-xsrf: 1234' localhost:5601/api/fleet/agent_policies/67c785b0-662e-11eb-bf6b-4790dc0178c0 -d '{ "name":"User created policy","namespace":"default","is_managed":true }'
{"item":{"id":"67c785b0-662e-11eb-bf6b-4790dc0178c0","name":"User created policy","namespace":"default","is_managed":true,"revision":2,"updated_at":"2021-02-03T14:47:28.471Z","updated_by":"elastic","package_policies":[]}}
```
</details>

## Enroll behavior
is not changed/addressed in this PR. Agents can still be enrolled in managed policies

## Unenroll Agent from managed policy behavior
#### Enrolled in managed agent policy, cannot be unenrolled
```
curl --user elastic:changeme -X POST http://localhost:5601/api/fleet/agents/441d4a40-6710-11eb-8f57-db14e8e41cff/unenroll -H 'kbn-xsrf: 1234' | jq
{
  "statusCode": 400,
  "error": "Bad Request",
  "message": "Cannot unenroll 441d4a40-6710-11eb-8f57-db14e8e41cff from a managed agent policy af9b4970-6701-11eb-b55a-899b78cb64da"
}
```

<details><summary>Screenshots for managed & unmanaged policies</summary>

#### Enrolled in managed agent policy, cannot be unenrolled
<img width="1931" alt="Screen Shot 2021-01-19 at 1 22 53 PM" src="https://user-images.githubusercontent.com/57655/105081614-67d05980-5a60-11eb-8faa-07e4e722a5b5.png">
<img width="1199" alt="Screen Shot 2021-01-19 at 1 30 26 PM" src="https://user-images.githubusercontent.com/57655/105081617-67d05980-5a60-11eb-9099-832dc6e04eca.png">
<img width="1971" alt="Screen Shot 2021-01-19 at 1 30 42 PM" src="https://user-images.githubusercontent.com/57655/105081618-67d05980-5a60-11eb-9a84-b80b6295ba19.png">

#### Enrolled agent policy is not managed, agent can be unenrolled<img width="1917" alt="Screen Shot 2021-01-19 at 1 44 12 PM" src="https://user-images.githubusercontent.com/57655/105081951-e3caa180-5a60-11eb-9308-7741b8986e8e.png">
<img width="2183" alt="Screen Shot 2021-01-19 at 1 44 19 PM" src="https://user-images.githubusercontent.com/57655/105081952-e3caa180-5a60-11eb-9833-1c721be0a107.png">

</details>


## Reassign agent 
#### No agent can be reassigned to a managed policy
```
 curl --user elastic:changeme -X 'PUT'  'http://localhost:5601/api/fleet/agents/482760d0-6710-11eb-8f57-db14e8e41cff/reassign' -H 'kbn-xsrf: xxx' -H 'Content-Type: application/json' -d '{"policy_id":"af9b4970-6701-11eb-b55a-899b78cb64da"}' 
{
  "statusCode": 400,
  "error": "Bad Request",
  "message": "Cannot reassign an agent to managed agent policy 94129590-6707-11eb-b55a-899b78cb64da"
}
```
<details><summary>Screenshots</summary>

<img width="1350" alt="Screen Shot 2021-02-04 at 2 14 51 PM" src="https://user-images.githubusercontent.com/57655/106943490-8044a300-66f3-11eb-9d2c-4b1ceef2e783.png">

</details>

#### Enrolled in managed agent policy, cannot be reassigned
```
 curl --user elastic:changeme -X 'PUT'  'http://localhost:5601/api/fleet/agents/482760d0-6710-11eb-8f57-db14e8e41cff/reassign' -H 'kbn-xsrf: xxx' -H 'Content-Type: application/json' -d '{"policy_id":"af9b4970-6701-11eb-b55a-899b78cb64da"}' 
{
  "statusCode": 400,
  "error": "Bad Request",
  "message": "Cannot reassign an agent from managed agent policy 94129590-6707-11eb-b55a-899b78cb64da"
}
```

<details><summary>Screenshots</summary>
<img width="1364" alt="Screen Shot 2021-01-19 at 2 58 38 PM" src="https://user-images.githubusercontent.com/57655/105086737-72dab800-5a67-11eb-8f5e-93cd7768b914.png">
<img width="1367" alt="Screen Shot 2021-01-19 at 2 58 44 PM" src="https://user-images.githubusercontent.com/57655/105086740-73734e80-5a67-11eb-8ef9-9c7005a0a4ea.png">
<img width="623" alt="Screen Shot 2021-01-19 at 2 59 27 PM" src="https://user-images.githubusercontent.com/57655/105086741-740be500-5a67-11eb-8fc2-721f8b5d178a.png">
</details>

#### Enrolled agent policy is unmanaged, agent can be reassigned to another unmanaged policy

<details><summary>Screenshots</summary>
<img width="1368" alt="Screen Shot 2021-01-19 at 3 00 01 PM" src="https://user-images.githubusercontent.com/57655/105086754-78d09900-5a67-11eb-86a5-9e3ac02d6e1f.png">
<img width="1363" alt="Screen Shot 2021-01-19 at 3 00 08 PM" src="https://user-images.githubusercontent.com/57655/105086761-7a01c600-5a67-11eb-991d-acf994e2a393.png">
<img width="625" alt="Screen Shot 2021-01-19 at 3 00 46 PM" src="https://user-images.githubusercontent.com/57655/105086764-7a9a5c80-5a67-11eb-8290-e79648d01579.png">
</details>

### Checklist

Delete any items that are not applicable to this PR.

- [ ] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/master/packages/kbn-i18n/README.md)
- [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials
- [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios

x-pack/plugins/fleet/common/constants/agent_policy.ts

@@ -24,6 +24,7 @@ export const DEFAULT_AGENT_POLICY: Omit<
   status: agentPolicyStatuses.Active,
   package_policies: [],
   is_default: true,
+  is_managed: false,
   monitoring_enabled: ['logs', 'metrics'] as Array<'logs' | 'metrics'>,
 };
 

x-pack/plugins/fleet/common/types/models/agent_policy.ts

@@ -17,13 +17,15 @@ export interface NewAgentPolicy {
   namespace: string;
   description?: string;
   is_default?: boolean;
+  is_managed?: boolean; // Optional when creating a policy
   monitoring_enabled?: Array<ValueOf<DataType>>;
 }
 
 export interface AgentPolicy extends NewAgentPolicy {
   id: string;
   status: ValueOf<AgentPolicyStatus>;
   package_policies: string[] | PackagePolicy[];
+  is_managed: boolean; // required for created policy
   updated_at: string;
   updated_by: string;
   revision: number;

x-pack/plugins/fleet/public/applications/fleet/sections/epm/screens/detail/index.test.tsx

@@ -687,6 +687,7 @@ On Windows, the module was tested with Nginx installed from the Chocolatey repos
           'e8a37031-2907-44f6-89d2-98bd493f60dc',
         ],
         is_default: true,
+        is_managed: false,
         monitoring_enabled: ['logs', 'metrics'],
         revision: 6,
         updated_at: '2020-12-09T13:46:31.840Z',
@@ -701,6 +702,7 @@ On Windows, the module was tested with Nginx installed from the Chocolatey repos
         status: 'active',
         package_policies: ['e8a37031-2907-44f6-89d2-98bd493f60cd'],
         is_default: false,
+        is_managed: false,
         monitoring_enabled: ['logs', 'metrics'],
         revision: 2,
         updated_at: '2020-12-09T13:46:31.840Z',

x-pack/plugins/fleet/server/errors/index.ts

@@ -32,3 +32,5 @@ export class PackageCacheError extends IngestManagerError {}
 export class PackageOperationNotSupportedError extends IngestManagerError {}
 export class FleetAdminUserInvalidError extends IngestManagerError {}
 export class ConcurrentInstallOperationError extends IngestManagerError {}
+export class AgentReassignmentError extends IngestManagerError {}
+export class AgentUnenrollmentError extends IngestManagerError {}

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

@@ -24,7 +24,7 @@ export const postAgentUnenrollHandler: RequestHandler<
     if (request.body?.force === true) {
       await AgentService.forceUnenrollAgent(soClient, esClient, request.params.agentId);
     } else {
-      await AgentService.unenrollAgent(soClient, request.params.agentId);
+      await AgentService.unenrollAgent(soClient, esClient, request.params.agentId);
     }
 
     const body: PostAgentUnenrollResponse = {};

x-pack/plugins/fleet/server/saved_objects/index.ts

@@ -32,7 +32,7 @@ import {
   migrateSettingsToV7100,
   migrateAgentActionToV7100,
 } from './migrations/to_v7_10_0';
-import { migrateAgentToV7120 } from './migrations/to_v7_12_0';
+import { migrateAgentToV7120, migrateAgentPolicyToV7120 } from './migrations/to_v7_12_0';
 
 /*
  * Saved object types and mappings
@@ -161,6 +161,7 @@ const getSavedObjectTypes = (
         description: { type: 'text' },
         namespace: { type: 'keyword' },
         is_default: { type: 'boolean' },
+        is_managed: { type: 'boolean' },
         status: { type: 'keyword' },
         package_policies: { type: 'keyword' },
         updated_at: { type: 'date' },
@@ -171,6 +172,7 @@ const getSavedObjectTypes = (
     },
     migrations: {
       '7.10.0': migrateAgentPolicyToV7100,
+      '7.12.0': migrateAgentPolicyToV7120,
     },
   },
   [ENROLLMENT_API_KEYS_SAVED_OBJECT_TYPE]: {

x-pack/plugins/fleet/server/saved_objects/migrations/to_v7_12_0.ts

@@ -5,8 +5,8 @@
  * 2.0.
  */
 
-import { SavedObjectMigrationFn } from 'kibana/server';
-import { Agent } from '../../types';
+import type { SavedObjectMigrationFn } from 'kibana/server';
+import type { Agent, AgentPolicy } from '../../types';
 
 export const migrateAgentToV7120: SavedObjectMigrationFn<Agent & { shared_id?: string }, Agent> = (
   agentDoc
@@ -15,3 +15,14 @@ export const migrateAgentToV7120: SavedObjectMigrationFn<Agent & { shared_id?: s
 
   return agentDoc;
 };
+
+export const migrateAgentPolicyToV7120: SavedObjectMigrationFn<
+  Exclude<AgentPolicy, 'is_managed'>,
+  AgentPolicy
+> = (agentPolicyDoc) => {
+  const isV12 = 'is_managed' in agentPolicyDoc.attributes;
+  if (!isV12) {
+    agentPolicyDoc.attributes.is_managed = false;
+  }
+  return agentPolicyDoc;
+};

x-pack/plugins/fleet/server/services/agent_policy.test.ts

@@ -8,17 +8,16 @@
 import { elasticsearchServiceMock, savedObjectsClientMock } from 'src/core/server/mocks';
 import { agentPolicyService } from './agent_policy';
 import { agentPolicyUpdateEventHandler } from './agent_policy_update';
-import { Output } from '../types';
+import type { AgentPolicy, NewAgentPolicy, Output } from '../types';
 
 function getSavedObjectMock(agentPolicyAttributes: any) {
   const mock = savedObjectsClientMock.create();
-
   mock.get.mockImplementation(async (type: string, id: string) => {
     return {
       type,
       id,
       references: [],
-      attributes: agentPolicyAttributes,
+      attributes: agentPolicyAttributes as AgentPolicy,
     };
   });
   mock.find.mockImplementation(async (options) => {
@@ -69,10 +68,59 @@ function getAgentPolicyUpdateMock() {
   >;
 }
 
+function getAgentPolicyCreateMock() {
+  const soClient = savedObjectsClientMock.create();
+  soClient.create.mockImplementation(async (type, attributes) => {
+    return {
+      attributes: (attributes as unknown) as NewAgentPolicy,
+      id: 'mocked',
+      type: 'mocked',
+      references: [],
+    };
+  });
+  return soClient;
+}
 describe('agent policy', () => {
   beforeEach(() => {
     getAgentPolicyUpdateMock().mockClear();
   });
+
+  describe('create', () => {
+    it('is_managed present and false by default', async () => {
+      // ignore unrelated unique name constraint
+      agentPolicyService.requireUniqueName = async () => {};
+      const soClient = getAgentPolicyCreateMock();
+      const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+
+      await expect(
+        agentPolicyService.create(soClient, esClient, {
+          name: 'No is_managed provided',
+          namespace: 'default',
+        })
+      ).resolves.toHaveProperty('is_managed', false);
+
+      const [, attributes] = soClient.create.mock.calls[0];
+      expect(attributes).toHaveProperty('is_managed', false);
+    });
+
+    it('should set is_managed property, if given', async () => {
+      // ignore unrelated unique name constraint
+      agentPolicyService.requireUniqueName = async () => {};
+      const soClient = getAgentPolicyCreateMock();
+      const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+      await expect(
+        agentPolicyService.create(soClient, esClient, {
+          name: 'is_managed: true provided',
+          namespace: 'default',
+          is_managed: true,
+        })
+      ).resolves.toHaveProperty('is_managed', true);
+
+      const [, attributes] = soClient.create.mock.calls[0];
+      expect(attributes).toHaveProperty('is_managed', true);
+    });
+  });
+
   describe('bumpRevision', () => {
     it('should call agentPolicyUpdateEventHandler with updated event once', async () => {
       const soClient = getSavedObjectMock({
@@ -208,4 +256,37 @@ describe('agent policy', () => {
       });
     });
   });
+
+  describe('update', () => {
+    it('should update is_managed property, if given', async () => {
+      // ignore unrelated unique name constraint
+      agentPolicyService.requireUniqueName = async () => {};
+      const soClient = savedObjectsClientMock.create();
+      const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+
+      soClient.get.mockResolvedValue({
+        attributes: {},
+        id: 'mocked',
+        type: 'mocked',
+        references: [],
+      });
+      await agentPolicyService.update(soClient, esClient, 'mocked', {
+        name: 'mocked',
+        namespace: 'default',
+        is_managed: false,
+      });
+      // soClient.update is called with updated values
+      let calledWith = soClient.update.mock.calls[0];
+      expect(calledWith[2]).toHaveProperty('is_managed', false);
+
+      await agentPolicyService.update(soClient, esClient, 'mocked', {
+        name: 'is_managed: true provided',
+        namespace: 'default',
+        is_managed: true,
+      });
+      // soClient.update is called with updated values
+      calledWith = soClient.update.mock.calls[1];
+      expect(calledWith[2]).toHaveProperty('is_managed', true);
+    });
+  });
 });

x-pack/plugins/fleet/server/services/agent_policy.ts

@@ -140,6 +140,7 @@ class AgentPolicyService {
       SAVED_OBJECT_TYPE,
       {
         ...agentPolicy,
+        is_managed: agentPolicy.is_managed ?? false,
         revision: 1,
         updated_at: new Date().toISOString(),
         updated_by: options?.user?.username || 'system',

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

@@ -12,7 +12,7 @@ import { AGENT_SAVED_OBJECT_TYPE } from '../../constants';
 import { AgentSOAttributes, Agent, ListWithKuery } from '../../types';
 import { escapeSearchQueryPhrase } from '../saved_object';
 import { savedObjectToAgent } from './saved_objects';
-import { appContextService } from '../../services';
+import { appContextService, agentPolicyService } from '../../services';
 import * as crudServiceSO from './crud_so';
 import * as crudServiceFleetServer from './crud_fleet_server';
 
@@ -86,6 +86,22 @@ export async function getAgents(soClient: SavedObjectsClientContract, agentIds:
   return agents;
 }
 
+export async function getAgentPolicyForAgent(
+  soClient: SavedObjectsClientContract,
+  esClient: ElasticsearchClient,
+  agentId: string
+) {
+  const agent = await getAgent(soClient, esClient, agentId);
+  if (!agent.policy_id) {
+    return;
+  }
+
+  const agentPolicy = await agentPolicyService.get(soClient, agent.policy_id, false);
+  if (agentPolicy) {
+    return agentPolicy;
+  }
+}
+
 export async function getAgentByAccessAPIKeyId(
   soClient: SavedObjectsClientContract,
   accessAPIKeyId: string

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

@@ -0,0 +1,132 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { elasticsearchServiceMock, savedObjectsClientMock } from 'src/core/server/mocks';
+import type { SavedObject } from 'kibana/server';
+import type { Agent, 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 unmanagedAgentPolicySO = {
+  id: 'unmanaged-agent-policy',
+  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);
+
+    // 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);
+  });
+
+  it('cannot reassign from unmanaged policy to managed', async () => {
+    const soClient = createClientMock();
+    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+    await expect(
+      reassignAgent(
+        soClient,
+        esClient,
+        agentInUnmanagedSO.id,
+        agentInManagedSO.attributes.policy_id!
+      )
+    ).rejects.toThrowError(AgentReassignmentError);
+
+    // does not call ES update
+    expect(soClient.update).toBeCalledTimes(0);
+  });
+
+  it('cannot reassign from managed policy', async () => {
+    const soClient = createClientMock();
+    const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+    await expect(
+      reassignAgent(soClient, esClient, agentInManagedSO.id, agentInManagedSO2.id)
+    ).rejects.toThrowError(AgentReassignmentError);
+    // does not call ES update
+    expect(soClient.update).toBeCalledTimes(0);
+
+    await expect(
+      reassignAgent(soClient, esClient, agentInManagedSO.id, agentInUnmanagedSO.id)
+    ).rejects.toThrowError(AgentReassignmentError);
+    // does not call ES update
+    expect(soClient.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);
+
+    // 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);
+  });
+});
+
+function createClientMock() {
+  const soClientMock = savedObjectsClientMock.create();
+
+  // need to mock .create & bulkCreate due to (bulk)createAgentAction(s) in reassignAgent(s)
+  soClientMock.create.mockResolvedValue(agentInUnmanagedSO);
+  soClientMock.bulkCreate.mockImplementation(async ([{ type, attributes }]) => {
+    return {
+      saved_objects: [await soClientMock.create(type, attributes)],
+    };
+  });
+
+  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:
+      default:
+        return agentInUnmanagedSO;
+    }
+  });
+
+  soClientMock.bulkGet.mockImplementation(async (options) => {
+    return {
+      saved_objects: await Promise.all(options!.map(({ type, id }) => soClientMock.get(type, id))),
+    };
+  });
+
+  return soClientMock;
+}