From 4c19c9c8b1587ac8097a0172ca5a0d32b315b911 Mon Sep 17 00:00:00 2001
From: Luke Heath <luke@fleetdm.com>
Date: Fri, 3 Dec 2021 10:13:40 -0600
Subject: [PATCH] Update policy back link to return to selected team (#3190)

---
 changes/issue-3130-fix-policy-back-link                    | 1 +
 .../policies/ManagePoliciesPage/ManagePoliciesPage.tsx     | 3 +--
 frontend/pages/policies/PolicyPage/screens/QueryEditor.tsx | 7 ++++++-
 frontend/pages/queries/QueryPage/screens/QueryEditor.tsx   | 2 +-
 frontend/pages/queries/QueryPage/screens/RunQuery.tsx      | 2 +-
 5 files changed, 10 insertions(+), 5 deletions(-)
 create mode 100644 changes/issue-3130-fix-policy-back-link

diff --git a/changes/issue-3130-fix-policy-back-link b/changes/issue-3130-fix-policy-back-link
new file mode 100644
index 0000000000..541e7110a7
--- /dev/null
+++ b/changes/issue-3130-fix-policy-back-link
@@ -0,0 +1 @@
+* Update policy back link to return to selected team
\ No newline at end of file
diff --git a/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx b/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
index cd319c865c..b08bdc4efb 100644
--- a/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
+++ b/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
@@ -1,5 +1,4 @@
 import React, { useCallback, useContext, useEffect, useState } from "react";
-import { Link } from "react-router";
 import { useQuery } from "react-query";
 import { useDispatch } from "react-redux";
 import { noop } from "lodash";
@@ -261,7 +260,7 @@ const ManagePolicyPage = (managePoliciesPageProps: {
     if (userTeams && !userTeams.find((t) => t.id === teamId)) {
       if (isOnGlobalTeam) {
         // For global users, default to zero (i.e. all teams).
-        if (teamId !== 0) {
+        if (teamId === undefined) {
           handleTeamSelect(0);
           return;
         }
diff --git a/frontend/pages/policies/PolicyPage/screens/QueryEditor.tsx b/frontend/pages/policies/PolicyPage/screens/QueryEditor.tsx
index a91358b259..853435467c 100644
--- a/frontend/pages/policies/PolicyPage/screens/QueryEditor.tsx
+++ b/frontend/pages/policies/PolicyPage/screens/QueryEditor.tsx
@@ -140,9 +140,14 @@ const QueryEditor = ({
     return null;
   }
 
+  const backPath = policyTeamId ? `?team_id=${policyTeamId}` : "";
+
   return (
     <div className={`${baseClass}__form body-wrap`}>
-      <Link to={PATHS.MANAGE_POLICIES} className={`${baseClass}__back-link`}>
+      <Link
+        to={`${PATHS.MANAGE_POLICIES}/${backPath}`}
+        className={`${baseClass}__back-link`}
+      >
         <img src={BackChevron} alt="back chevron" id="back-chevron" />
         <span>Back to policies</span>
       </Link>
diff --git a/frontend/pages/queries/QueryPage/screens/QueryEditor.tsx b/frontend/pages/queries/QueryPage/screens/QueryEditor.tsx
index 7d55a6f491..b94b7911da 100644
--- a/frontend/pages/queries/QueryPage/screens/QueryEditor.tsx
+++ b/frontend/pages/queries/QueryPage/screens/QueryEditor.tsx
@@ -44,7 +44,7 @@ const QueryEditor = ({
   goToSelectTargets,
   onOpenSchemaSidebar,
   renderLiveQueryWarning,
-}: IQueryEditorProps) => {
+}: IQueryEditorProps): JSX.Element | null => {
   const dispatch = useDispatch();
   const { currentUser } = useContext(AppContext);
 
diff --git a/frontend/pages/queries/QueryPage/screens/RunQuery.tsx b/frontend/pages/queries/QueryPage/screens/RunQuery.tsx
index ff55c50967..45e04fae9e 100644
--- a/frontend/pages/queries/QueryPage/screens/RunQuery.tsx
+++ b/frontend/pages/queries/QueryPage/screens/RunQuery.tsx
@@ -32,7 +32,7 @@ const RunQuery = ({
   queryIdForEdit,
   setSelectedTargets,
   goToQueryEditor,
-}: IRunQueryProps) => {
+}: IRunQueryProps): JSX.Element | null => {
   const dispatch = useDispatch();
 
   const [isQueryFinished, setIsQueryFinished] = useState<boolean>(false);