From dc8b07d46ffad945715df6035f48ecd6a3f5f723 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=A1bor=20Farkas?= <gabor.farkas@gmail.com>
Date: Wed, 11 Oct 2023 10:38:38 +0200
Subject: [PATCH] postgres/mysql: improved imports in unit tests (#75980)

* postgres/mysql: improved imports in unit tests

* simpler mock approach
---
 .../datasource/mysql/specs/datasource.test.ts | 13 ++++--
 .../datasource/postgres/datasource.test.ts    | 46 +++++++++++++------
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/public/app/plugins/datasource/mysql/specs/datasource.test.ts b/public/app/plugins/datasource/mysql/specs/datasource.test.ts
index 335823e8594..aaee5eabf0b 100644
--- a/public/app/plugins/datasource/mysql/specs/datasource.test.ts
+++ b/public/app/plugins/datasource/mysql/specs/datasource.test.ts
@@ -9,20 +9,25 @@ import {
   FieldType,
   createDataFrame,
 } from '@grafana/data';
-import { FetchResponse, setBackendSrv } from '@grafana/runtime';
-import { backendSrv } from 'app/core/services/backend_srv'; // will use the version in __mocks__
+import { FetchResponse } from '@grafana/runtime';
 import { SQLQuery } from 'app/features/plugins/sql/types';
 import { makeVariable } from 'app/features/plugins/sql/utils/testHelpers';
 
 import { MySqlDatasource } from '../MySqlDatasource';
 import { MySQLOptions } from '../types';
 
+const fetchMock = jest.fn();
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  getBackendSrv: () => ({
+    fetch: fetchMock,
+  }),
+}));
+
 describe('MySQLDatasource', () => {
   const defaultRange = getDefaultTimeRange(); // it does not matter what value this has
   const setupTestContext = (response: unknown, templateSrv?: unknown) => {
     jest.clearAllMocks();
-    setBackendSrv(backendSrv);
-    const fetchMock = jest.spyOn(backendSrv, 'fetch');
     const instanceSettings = {
       jsonData: {
         defaultProject: 'testproject',
diff --git a/public/app/plugins/datasource/postgres/datasource.test.ts b/public/app/plugins/datasource/postgres/datasource.test.ts
index da814d026a9..f5379945829 100644
--- a/public/app/plugins/datasource/postgres/datasource.test.ts
+++ b/public/app/plugins/datasource/postgres/datasource.test.ts
@@ -12,28 +12,44 @@ import {
   LoadingState,
   createDataFrame,
 } from '@grafana/data';
-import { FetchResponse } from '@grafana/runtime';
-import { backendSrv } from 'app/core/services/backend_srv'; // will use the version in __mocks__
+import {
+  BackendSrv,
+  DataSourceSrv,
+  FetchResponse,
+  getBackendSrv,
+  setBackendSrv,
+  getDataSourceSrv,
+  setDataSourceSrv,
+} from '@grafana/runtime';
 import { QueryFormat, SQLQuery } from 'app/features/plugins/sql/types';
 import { makeVariable } from 'app/features/plugins/sql/utils/testHelpers';
 
 import { PostgresDatasource } from './datasource';
 import { PostgresOptions } from './types';
 
-jest.mock('@grafana/runtime', () => ({
-  ...jest.requireActual('@grafana/runtime'),
-  getBackendSrv: () => backendSrv,
-}));
+const backendSrv: BackendSrv = {
+  // this will get mocked below, it only needs to exist
+  fetch: () => undefined,
+} as unknown as BackendSrv; // we cast it so that we do not have to implement all the methods
 
-jest.mock('@grafana/runtime/src/services', () => ({
-  ...jest.requireActual('@grafana/runtime/src/services'),
-  getBackendSrv: () => backendSrv,
-  getDataSourceSrv: () => {
-    return {
-      getInstanceSettings: () => ({ id: 8674 }),
-    };
-  },
-}));
+// we type this as `any` to not have to define the whole type
+const fakeDataSourceSrv: DataSourceSrv = {
+  getInstanceSettings: () => ({ id: 8674 }),
+} as unknown as DataSourceSrv;
+
+let origBackendSrv: BackendSrv;
+let origDataSourceSrv: DataSourceSrv;
+beforeAll(() => {
+  origBackendSrv = getBackendSrv();
+  origDataSourceSrv = getDataSourceSrv();
+  setBackendSrv(backendSrv);
+  setDataSourceSrv(fakeDataSourceSrv);
+});
+
+afterAll(() => {
+  setBackendSrv(origBackendSrv);
+  setDataSourceSrv(origDataSourceSrv);
+});
 
 describe('PostgreSQLDatasource', () => {
   const defaultRange = getDefaultTimeRange(); // it does not matter what value this has