From 10e66360a84e14f16c6577c9d48fa6062583b9c9 Mon Sep 17 00:00:00 2001
From: Shenoy Pratik <sgguruda@amazon.com>
Date: Mon, 4 Nov 2024 08:15:18 -0800
Subject: [PATCH] Fix mds ref update on integration assets (#2240)

* fix mds ref update on integration assests

Signed-off-by: Shenoy Pratik <sgguruda@amazon.com>

* add unit test for builder with mds enabled

Signed-off-by: Shenoy Pratik <sgguruda@amazon.com>

---------

Signed-off-by: Shenoy Pratik <sgguruda@amazon.com>
---
 .../components/create_integration_helpers.ts  |  3 +-
 .../integrations/__test__/builder.test.ts     | 89 +++++++++++++++++++
 .../integrations/integrations_builder.ts      |  2 +-
 3 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/public/components/integrations/components/create_integration_helpers.ts b/public/components/integrations/components/create_integration_helpers.ts
index f6b8ed4c8a..7b05b8332a 100644
--- a/public/components/integrations/components/create_integration_helpers.ts
+++ b/public/components/integrations/components/create_integration_helpers.ts
@@ -305,7 +305,7 @@ export async function addIntegrationRequest({
   indexPattern,
   workflows,
   skipRedirect,
-  dataSourceInfo,
+  dataSourceInfo, // s3 datasource
   dataSourceMDSId,
   dataSourceMDSLabel,
 }: AddIntegrationRequestParams): Promise<boolean> {
@@ -378,6 +378,7 @@ export async function addIntegrationRequest({
       query: {
         path: `${indexPattern}/_bulk?refresh=wait_for`,
         method: 'POST',
+        dataSourceId: dataSourceMDSId,
       },
     })
     .then((_) => {
diff --git a/server/adaptors/integrations/__test__/builder.test.ts b/server/adaptors/integrations/__test__/builder.test.ts
index 4ebec0ad35..b685993890 100644
--- a/server/adaptors/integrations/__test__/builder.test.ts
+++ b/server/adaptors/integrations/__test__/builder.test.ts
@@ -129,6 +129,95 @@ describe('IntegrationInstanceBuilder', () => {
       expect(instance).toEqual(expectedInstance);
     });
 
+    it('should build an integration instance with MDS reference', async () => {
+      const options = {
+        indexPattern: 'instance-datasource',
+        name: 'instance-name',
+        dataSourceMDSId: 'sampleMDSId',
+        dataSourceMDSLabel: 'sample MDS',
+      };
+
+      const remappedAssets = [
+        {
+          id: 'remapped-asset1',
+          references: [{ id: 'remapped-ref1' }],
+        },
+        {
+          id: 'remapped-asset2',
+          references: [{ id: 'remapped-ref2' }],
+        },
+      ];
+      const postAssetsResponse = {
+        saved_objects: [
+          { id: 'created-asset1', type: 'dashboard', attributes: { title: 'Dashboard 1' } },
+          { id: 'created-asset2', type: 'visualization', attributes: { title: 'Visualization 1' } },
+        ],
+      };
+      const expectedInstance = {
+        name: 'instance-name',
+        templateName: 'sample',
+        dataSource: 'instance-datasource',
+        creationDate: expect.any(String),
+        assets: [
+          {
+            assetType: 'dashboard',
+            assetId: 'created-asset1',
+            status: 'available',
+            isDefaultAsset: true,
+            description: 'Dashboard 1',
+          },
+          {
+            assetType: 'visualization',
+            assetId: 'created-asset2',
+            status: 'available',
+            isDefaultAsset: false,
+            description: 'Visualization 1',
+          },
+        ],
+        references: [
+          {
+            id: 'sampleMDSId',
+            name: 'sample MDS',
+            type: 'data-source',
+          },
+        ],
+      };
+
+      const mockTemplate: Partial<IntegrationConfig> = TEST_INTEGRATION_CONFIG;
+
+      jest
+        .spyOn(mockUtils, 'deepCheck')
+        .mockResolvedValue({ ok: true, value: mockTemplate as IntegrationConfig });
+
+      // Mock the implementation of the methods in the Integration class
+      // sampleIntegration.deepCheck = jest.fn().mockResolvedValue({ ok: true, value: mockTemplate });
+      sampleIntegration.getAssets = jest.fn().mockResolvedValue({
+        ok: true,
+        value: [{ type: 'savedObjectBundle', data: remappedAssets }],
+      });
+      sampleIntegration.getConfig = jest.fn().mockResolvedValue({ ok: true, value: mockTemplate });
+
+      // Mock builder sub-methods
+      const remapIDsSpy = jest.spyOn(builder, 'remapIDs');
+      const addMDSReferenceSpy = jest.spyOn(builder, 'addMDSReference');
+      const postAssetsSpy = jest.spyOn(builder, 'postAssets');
+
+      (mockSavedObjectsClient.bulkCreate as jest.Mock).mockResolvedValue(postAssetsResponse);
+
+      const instance = await builder.build(sampleIntegration, options);
+
+      expect(sampleIntegration.getAssets).toHaveBeenCalled();
+      expect(remapIDsSpy).toHaveBeenCalledWith(remappedAssets);
+      expect(addMDSReferenceSpy).toHaveBeenCalledWith(
+        remappedAssets,
+        undefined,
+        'sampleMDSId',
+        'sample MDS'
+      );
+      expect(postAssetsSpy).toHaveBeenCalledWith(remappedAssets);
+      expect(instance).toEqual(expectedInstance);
+    });
+
     it('should reject with an error if integration is not valid', async () => {
       const options = {
         indexPattern: 'instance-datasource',
diff --git a/server/adaptors/integrations/integrations_builder.ts b/server/adaptors/integrations/integrations_builder.ts
index c908d2e2fe..e709cfaab9 100644
--- a/server/adaptors/integrations/integrations_builder.ts
+++ b/server/adaptors/integrations/integrations_builder.ts
@@ -207,7 +207,7 @@ export class IntegrationInstanceBuilder {
     dataSourceMDSId?: string,
     dataSourceMDSLabel?: string
   ): SavedObject[] {
-    if (!dataSource) {
+    if (!dataSourceMDSId) {
       return assets;
     }
     return assets.map((asset) => {