From c8f3934772bf2766f7a6d1d151fcd9a174dd06c9 Mon Sep 17 00:00:00 2001
From: gillespi314 <73313222+gillespi314@users.noreply.github.com>
Date: Wed, 17 Nov 2021 10:26:49 -0600
Subject: [PATCH] Fix targets selector bug (#2981)

---
 changes/issue-2953-select-targets             |  1 +
 .../QueryPage/screens/SelectTargets.tsx       | 91 ++++++++++++-------
 2 files changed, 60 insertions(+), 32 deletions(-)
 create mode 100644 changes/issue-2953-select-targets

diff --git a/changes/issue-2953-select-targets b/changes/issue-2953-select-targets
new file mode 100644
index 0000000000..f1e0d9a996
--- /dev/null
+++ b/changes/issue-2953-select-targets
@@ -0,0 +1 @@
+* Fix target selector bug in case where different target types have the same numeric id
\ No newline at end of file
diff --git a/frontend/pages/queries/QueryPage/screens/SelectTargets.tsx b/frontend/pages/queries/QueryPage/screens/SelectTargets.tsx
index 5a21f9ea50..bf3d180692 100644
--- a/frontend/pages/queries/QueryPage/screens/SelectTargets.tsx
+++ b/frontend/pages/queries/QueryPage/screens/SelectTargets.tsx
@@ -1,7 +1,7 @@
 import React, { useState } from "react";
 import { useQuery } from "react-query";
 import { Row } from "react-table";
-import { filter, forEach, isEmpty, remove, unionBy } from "lodash";
+import { filter, forEach, isEmpty, remove, unionWith } from "lodash";
 
 // @ts-ignore
 import { formatSelectedTargetsForApi } from "fleet/helpers";
@@ -20,17 +20,31 @@ import CheckIcon from "../../../../../assets/images/icon-check-purple-32x32@2x.p
 import ExternalURLIcon from "../../../../../assets/images/icon-external-url-12x12@2x.png";
 import ErrorIcon from "../../../../../assets/images/icon-error-16x16@2x.png";
 
+interface ISelectHost extends IHost {
+  target_type?: string;
+}
+
+interface ISelectLabel extends ILabel {
+  target_type?: string;
+}
+
+interface ISelectTeam extends ITeam {
+  target_type?: string;
+}
+
+type ISelectTargetsEntity = ISelectHost | ISelectLabel | ISelectTeam;
+
 interface ITargetPillSelectorProps {
-  entity: ILabel | ITeam;
+  entity: ISelectLabel | ISelectTeam;
   isSelected: boolean;
   onClick: (
-    value: ILabel | ITeam
+    value: ISelectLabel | ISelectTeam
   ) => React.MouseEventHandler<HTMLButtonElement>;
 }
 
 interface ISelectTargetsProps {
   baseClass: string;
-  selectedTargets: ITarget[];
+  selectedTargets: ISelectTargetsEntity[];
   queryIdForEdit: number | null;
   goToQueryEditor: () => void;
   goToRunQuery: () => void;
@@ -43,6 +57,11 @@ interface IModifiedUseQueryTargetsResponse {
   onlineCount: number;
 }
 
+const isSameSelectTargetsEntity = (
+  e1: ISelectTargetsEntity,
+  e2: ISelectTargetsEntity
+) => e1.id === e2.id && e1.target_type === e2.target_type;
+
 const TargetPillSelector = ({
   entity,
   isSelected,
@@ -88,7 +107,9 @@ const SelectTargets = ({
   const [platformLabels, setPlatformLabels] = useState<ILabel[] | null>(null);
   const [teams, setTeams] = useState<ITeam[] | null>(null);
   const [otherLabels, setOtherLabels] = useState<ILabel[] | null>(null);
-  const [selectedLabels, setSelectedLabels] = useState<any>([]);
+  const [selectedLabels, setSelectedLabels] = useState<ISelectTargetsEntity[]>(
+    []
+  );
   const [inputTabIndex, setInputTabIndex] = useState<number>(0);
   const [searchText, setSearchText] = useState<string>("");
   const [relatedHosts, setRelatedHosts] = useState<IHost[]>([]);
@@ -169,32 +190,36 @@ const SelectTargets = ({
     }
   );
 
-  const handleSelectedLabels = (entity: ILabel | ITeam) => (
+  const handleSelectedLabels = (selectedLabel: ISelectTargetsEntity) => (
     e: React.MouseEvent<HTMLButtonElement>
   ): void => {
     e.preventDefault();
+
+    let targets = selectedTargets;
     const labels = selectedLabels;
-    let newTargets = null;
-    const targets = selectedTargets;
-    const removed = remove(labels, ({ id }) => id === entity.id);
+    const removed = remove(labels, (label) =>
+      isSameSelectTargetsEntity(label, selectedLabel)
+    );
 
     // visually show selection
     const isRemoval = removed.length > 0;
     if (isRemoval) {
-      newTargets = labels;
+      targets = targets.filter(
+        (target) => !isSameSelectTargetsEntity(target, selectedLabel)
+      );
     } else {
-      labels.push(entity);
+      labels.push(selectedLabel);
 
       // prepare the labels data
       forEach(labels, (label) => {
         label.target_type = "label_type" in label ? "labels" : "teams";
       });
 
-      newTargets = unionBy(targets, labels, "id");
+      targets = unionWith(targets, labels, isSameSelectTargetsEntity);
     }
 
     setSelectedLabels([...labels]);
-    setSelectedTargets([...newTargets]);
+    setSelectedTargets([...targets]);
   };
 
   const handleRowSelect = (row: Row) => {
@@ -211,31 +236,33 @@ const SelectTargets = ({
   const handleRowRemove = (row: Row) => {
     const targets = selectedTargets;
     const hostTarget = row.original as ITarget;
-    remove(targets, (t) => t.id === hostTarget.id);
+    remove(targets, (t) => t.id === hostTarget.id && t.target_type === "hosts");
 
     setSelectedTargets([...targets]);
   };
 
   const renderTargetEntityList = (
     header: string,
-    entityList: ILabel[] | ITeam[]
-  ): JSX.Element => (
-    <>
-      {header && <h3>{header}</h3>}
-      <div className="selector-block">
-        {entityList?.map((entity: ILabel | ITeam) => (
-          <TargetPillSelector
-            key={entity.id}
-            entity={entity}
-            isSelected={selectedLabels.some(
-              ({ id }: ILabel | ITeam) => id === entity.id
-            )}
-            onClick={handleSelectedLabels}
-          />
-        ))}
-      </div>
-    </>
-  );
+    entityList: ISelectLabel[] | ISelectTeam[]
+  ): JSX.Element => {
+    return (
+      <>
+        {header && <h3>{header}</h3>}
+        <div className="selector-block">
+          {entityList?.map((entity: ISelectLabel | ISelectTeam) => (
+            <TargetPillSelector
+              key={`target_${entity.target_type}_${entity.id}`}
+              entity={entity}
+              isSelected={selectedLabels.some((label: ISelectTargetsEntity) =>
+                isSameSelectTargetsEntity(label, entity)
+              )}
+              onClick={handleSelectedLabels}
+            />
+          ))}
+        </div>
+      </>
+    );
+  };
 
   if (isEmpty(searchText) && isTargetsLoading) {
     return (