From 9ec8db5681d773e99110b8bc590b15a9800a43fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?H=C3=A5kon=20Flatval?= <hakon.flatval@cognite.com>
Date: Tue, 1 Aug 2023 11:11:15 +0200
Subject: [PATCH] chore: refactor

---
 react-components/stories/HighlightNode.stories.tsx | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/react-components/stories/HighlightNode.stories.tsx b/react-components/stories/HighlightNode.stories.tsx
index 61558eb6ac0..0f92827980b 100644
--- a/react-components/stories/HighlightNode.stories.tsx
+++ b/react-components/stories/HighlightNode.stories.tsx
@@ -78,16 +78,12 @@ const StoryContent = ({
   resources: AddResourceOptions[];
   fdmAssetMappingConfig: FdmAssetMappingsConfig;
 }) => {
-  const viewer = useReveal();
-
   const [nodeData, setNodeData] = useState<any>();
 
   const callback = (nodeData: NodeDataResult<any>) => {
     setNodeData(nodeData.data);
 
-    if (!(viewer.models[0] instanceof CogniteCadModel)) return;
-
-    viewer.models[0].assignStyledNodeCollection(
+    nodeData.model.assignStyledNodeCollection(
       new TreeIndexNodeCollection([nodeData.cadNode.treeIndex]),
       DefaultNodeAppearance.Highlighted
     );