From 74c1c21942509315469178f954298271c775de27 Mon Sep 17 00:00:00 2001
From: Yevhenii Huselietov <yevhenii.huselietov@gmail.com>
Date: Thu, 5 Oct 2023 14:34:30 +0300
Subject: [PATCH] Review comments

Signed-off-by: Yevhenii Huselietov <yevhenii.huselietov@gmail.com>
---
 .../NewRelicDashboard.stories.tsx             | 47 ++++++++++---------
 1 file changed, 24 insertions(+), 23 deletions(-)

diff --git a/plugins/newrelic-dashboard/src/components/NewRelicDashboard/NewRelicDashboard.stories.tsx b/plugins/newrelic-dashboard/src/components/NewRelicDashboard/NewRelicDashboard.stories.tsx
index c4f81b8321..a15a4c573c 100644
--- a/plugins/newrelic-dashboard/src/components/NewRelicDashboard/NewRelicDashboard.stories.tsx
+++ b/plugins/newrelic-dashboard/src/components/NewRelicDashboard/NewRelicDashboard.stories.tsx
@@ -21,6 +21,7 @@ import { EntityProvider } from '@backstage/plugin-catalog-react';
 import { newRelicDashboardApiRef } from '../../api';
 import { NEWRELIC_GUID_ANNOTATION } from '../../constants';
 import { storageApiRef } from '@backstage/core-plugin-api';
+import { Entity } from '@backstage/catalog-model';
 
 function createImage(
   width: number,
@@ -44,13 +45,15 @@ function createImage(
   return canvas;
 }
 
-const entity: any = {
+const entity: Entity = {
   apiVersion: '1',
   kind: 'Component',
   metadata: {
     name: 'mocked entity with newrelic service',
     title: 'app with newrelic',
-    [NEWRELIC_GUID_ANNOTATION]: 'some-cool-guid',
+    annotations: {
+      [NEWRELIC_GUID_ANNOTATION]: 'some-cool-guid',
+    },
   },
 };
 
@@ -58,21 +61,13 @@ const newRelicApiMockEmpty = {
   getDashboardEntity: () => Promise.resolve(undefined),
 };
 
-const mockedNewRelicDashboard = (apis: any[]) => {
-  return (
-    <TestApiProvider apis={apis}>
-      <EntityProvider entity={entity}>
-        <NewRelicDashboard />
-      </EntityProvider>
-    </TestApiProvider>
-  );
-};
-
-export const EmptyNewRelicDashboard = () => {
-  return mockedNewRelicDashboard([
-    [newRelicDashboardApiRef, newRelicApiMockEmpty],
-  ]);
-};
+export const EmptyNewRelicDashboard = () => (
+  <TestApiProvider apis={[[newRelicDashboardApiRef, newRelicApiMockEmpty]]}>
+    <EntityProvider entity={entity}>
+      <NewRelicDashboard />
+    </EntityProvider>
+  </TestApiProvider>
+);
 
 const resultEntities = [
   {
@@ -117,12 +112,18 @@ const newRelicApiMockFull = {
   getDashboardSnapshot: () => Promise.resolve(dashboardSnapshot),
 };
 
-export const NewRelicDashboardWithSnapshots = () => {
-  return mockedNewRelicDashboard([
-    [newRelicDashboardApiRef, newRelicApiMockFull],
-    [storageApiRef, MockStorageApi.create()],
-  ]);
-};
+export const NewRelicDashboardWithSnapshots = () => (
+  <TestApiProvider
+    apis={[
+      [newRelicDashboardApiRef, newRelicApiMockFull],
+      [storageApiRef, MockStorageApi.create()],
+    ]}
+  >
+    <EntityProvider entity={entity}>
+      <NewRelicDashboard />
+    </EntityProvider>
+  </TestApiProvider>
+);
 
 export default {
   title: 'NewRelic Dashboard',