From f5958fe0c0c67e8c05211719be7350244ea21f85 Mon Sep 17 00:00:00 2001
From: Brendan O'Handley <brendan.ohandley@grafana.com>
Date: Mon, 25 Nov 2024 16:36:58 -0600
Subject: [PATCH] Explore metrics: Fix default datasource bug, add test
 coverage (#97008)

* select actual default prom data source

* add tests for the function
---
 public/app/features/trails/utils.test.ts | 58 +++++++++++++++++++++++-
 public/app/features/trails/utils.ts      |  4 +-
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/public/app/features/trails/utils.test.ts b/public/app/features/trails/utils.test.ts
index e360d1c4da3..f4cab00bc54 100644
--- a/public/app/features/trails/utils.test.ts
+++ b/public/app/features/trails/utils.test.ts
@@ -1,8 +1,19 @@
-import { AdHocFiltersVariable } from '@grafana/scenes';
+import { AdHocFiltersVariable, SceneObjectRef } from '@grafana/scenes';
+
+import { getDatasourceSrv } from '../plugins/datasource_srv';
 
 import { DataTrail } from './DataTrail';
+import { getTrailStore } from './TrailStore/TrailStore';
 import { MetricDatasourceHelper } from './helpers/MetricDatasourceHelper';
-import { limitAdhocProviders } from './utils';
+import { getDatasourceForNewTrail, limitAdhocProviders } from './utils';
+
+jest.mock('./TrailStore/TrailStore', () => ({
+  getTrailStore: jest.fn(),
+}));
+
+jest.mock('../plugins/datasource_srv', () => ({
+  getDatasourceSrv: jest.fn(),
+}));
 
 describe('limitAdhocProviders', () => {
   let filtersVariable: AdHocFiltersVariable;
@@ -60,3 +71,46 @@ describe('limitAdhocProviders', () => {
     }
   });
 });
+
+describe('getDatasourceForNewTrail', () => {
+  beforeEach(() => {
+    (getTrailStore as jest.Mock).mockImplementation(() => ({
+      bookmarks: [],
+      recent: [],
+    }));
+    (getDatasourceSrv as jest.Mock).mockImplementation(() => ({
+      getList: jest.fn().mockReturnValue([
+        { uid: 'prom1', isDefault: true },
+        { uid: 'prom2', isDefault: false },
+      ]),
+    }));
+  });
+
+  it('should return the most recent exploration data source', () => {
+    const trail = new DataTrail({ key: '1', metric: 'select me', initialDS: 'prom2' });
+    const trailWithResolveMethod = new SceneObjectRef(trail);
+    (getTrailStore as jest.Mock).mockImplementation(() => ({
+      bookmarks: [],
+      recent: [trailWithResolveMethod],
+    }));
+    const result = getDatasourceForNewTrail();
+    expect(result).toBe('prom2');
+  });
+
+  it('should return the default Prometheus data source if no previous exploration exists', () => {
+    const result = getDatasourceForNewTrail();
+    expect(result).toBe('prom1');
+  });
+
+  it('should return the most recently added Prom data source if no default exists and no recent exploration', () => {
+    (getDatasourceSrv as jest.Mock).mockImplementation(() => ({
+      getList: jest.fn().mockReturnValue([
+        { uid: 'newProm', isDefault: false },
+        { uid: 'prom1', isDefault: false },
+        { uid: 'prom2', isDefault: false },
+      ]),
+    }));
+    const result = getDatasourceForNewTrail();
+    expect(result).toBe('newProm');
+  });
+});
diff --git a/public/app/features/trails/utils.ts b/public/app/features/trails/utils.ts
index 317271dfbd6..4045fffba3d 100644
--- a/public/app/features/trails/utils.ts
+++ b/public/app/features/trails/utils.ts
@@ -103,7 +103,9 @@ export function getDatasourceForNewTrail(): string | undefined {
   }
   const promDatasources = getDatasourceSrv().getList({ type: 'prometheus' });
   if (promDatasources.length > 0) {
-    return promDatasources.find((mds) => mds.uid === config.defaultDatasource)?.uid ?? promDatasources[0].uid;
+    const defaultDatasource = promDatasources.find((mds) => mds.isDefault);
+
+    return defaultDatasource?.uid ?? promDatasources[0].uid;
   }
   return undefined;
 }