[Discover] Some cleanups for the new in-table search (elastic#208980)

- Addresses elastic#208939

## Summary

This PR makes some cleanups to the code introduced in
elastic#206454 and adds more tests.

### Checklist

- [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

(cherry picked from commit b1b28c3)

src/platform/packages/shared/kbn-data-grid-in-table-search/src/__mocks__/index.ts

@@ -10,3 +10,4 @@
 export { generateMockData } from './data';
 export { getRenderCellValueMock } from './render_cell_value_mock';
 export { DataGridWithInTableSearchExample } from './data_grid_example';
+export { MockContext, useMockContextValue } from './mock_context';

src/platform/packages/shared/kbn-data-grid-in-table-search/src/__mocks__/mock_context.ts

@@ -0,0 +1,18 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import React from 'react';
+
+interface MockContextValue {
+  mockContextValue?: string;
+}
+
+export const MockContext = React.createContext<MockContextValue>({});
+
+export const useMockContextValue = () => React.useContext(MockContext).mockContextValue;

src/platform/packages/shared/kbn-data-grid-in-table-search/src/__mocks__/render_cell_value_mock.tsx

@@ -9,18 +9,28 @@
 
 import React from 'react';
 import type { EuiDataGridCellValueElementProps } from '@elastic/eui';
+import { useMockContextValue } from './mock_context';
 
 export function getRenderCellValueMock(testData: string[][]) {
   return function OriginalRenderCellValue({
     colIndex,
     rowIndex,
   }: EuiDataGridCellValueElementProps) {
+    const mockContextValue = useMockContextValue();
     const cellValue = testData[rowIndex][colIndex];
 
     if (!cellValue) {
       throw new Error('Testing unexpected errors');
     }
 
-    return <div>{cellValue}</div>;
+    return (
+      <div>
+        {cellValue}
+        {
+          // testing that it can access the parent context value
+          mockContextValue ? <span>{mockContextValue}</span> : null
+        }
+      </div>
+    );
   };
 }