diff --git a/public/contexts/__mocks__/core_context.tsx b/public/contexts/__mocks__/core_context.tsx
index 5e8770ba..2d92d77f 100644
--- a/public/contexts/__mocks__/core_context.tsx
+++ b/public/contexts/__mocks__/core_context.tsx
@@ -5,6 +5,7 @@
 
 import { BehaviorSubject } from 'rxjs';
 import { coreMock } from '../../../../../src/core/public/mocks';
+import { DataSourceServiceMock } from '../../services/data_source_service.mock';
 
 export const useCore = jest.fn(() => {
   const useCoreMock = {
@@ -24,11 +25,7 @@ export const useCore = jest.fn(() => {
         load: jest.fn(),
       },
       conversationLoad: {},
-      dataSource: {
-        getDataSourceQuery() {
-          return { dataSourceId: '' };
-        },
-      },
+      dataSource: new DataSourceServiceMock(),
     },
   };
   useCoreMock.services.http.delete.mockReturnValue(Promise.resolve());