From ced40d5aeced41b02f44cab7e8329d6d39ca3345 Mon Sep 17 00:00:00 2001
From: Harman-singh-waraich <Waraichharman68@gmail.com>
Date: Wed, 7 Aug 2024 19:59:52 +0530
Subject: [PATCH] refactor(web): improve-draw-maintenance-button

---
 .../MaintenanceButtons/DrawButton.tsx         | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/web/src/pages/Cases/CaseDetails/MaintenanceButtons/DrawButton.tsx b/web/src/pages/Cases/CaseDetails/MaintenanceButtons/DrawButton.tsx
index f4a23a654..67b4a77e0 100644
--- a/web/src/pages/Cases/CaseDetails/MaintenanceButtons/DrawButton.tsx
+++ b/web/src/pages/Cases/CaseDetails/MaintenanceButtons/DrawButton.tsx
@@ -6,6 +6,7 @@ import { usePublicClient } from "wagmi";
 import { Button } from "@kleros/ui-components-library";
 
 import { useSimulateKlerosCoreDraw, useWriteKlerosCoreDraw } from "hooks/contracts/generated";
+import { useSortitionModulePhase } from "hooks/useSortitionModulePhase";
 import { wrapWithToast } from "utils/wrapWithToast";
 
 import useDisputeMaintenanceQuery from "queries/useDisputeMaintenanceQuery";
@@ -13,6 +14,8 @@ import useDisputeMaintenanceQuery from "queries/useDisputeMaintenanceQuery";
 import { Period } from "src/graphql/graphql";
 import { isUndefined } from "src/utils";
 
+import { Phases } from "components/Phase";
+
 import { IBaseMaintenanceButton } from ".";
 
 const StyledButton = styled(Button)`
@@ -28,21 +31,22 @@ const DrawButton: React.FC<IDrawButton> = ({ id, numberOfVotes, setIsOpen, perio
   const [isSending, setIsSending] = useState(false);
   const publicClient = usePublicClient();
   const { data: maintenanceData } = useDisputeMaintenanceQuery(id);
+  const { data: phase } = useSortitionModulePhase();
 
   const isDrawn = useMemo(() => maintenanceData?.dispute?.currentRound.jurorsDrawn, [maintenanceData]);
 
+  const canDraw = useMemo(
+    () => !isUndefined(maintenanceData) && !isDrawn && period === Period.Evidence && phase === Phases.drawing,
+    [maintenanceData, isDrawn, phase, period]
+  );
+
   const {
     data: drawConfig,
     isLoading: isLoadingConfig,
     isError,
   } = useSimulateKlerosCoreDraw({
     query: {
-      enabled:
-        !isUndefined(id) &&
-        !isUndefined(numberOfVotes) &&
-        !isUndefined(period) &&
-        period === Period.Evidence &&
-        !isDrawn,
+      enabled: !isUndefined(id) && !isUndefined(numberOfVotes) && !isUndefined(period) && canDraw,
     },
     args: [BigInt(id ?? 0), BigInt(numberOfVotes ?? 0)],
   });
@@ -51,9 +55,8 @@ const DrawButton: React.FC<IDrawButton> = ({ id, numberOfVotes, setIsOpen, perio
 
   const isLoading = useMemo(() => isLoadingConfig || isSending, [isLoadingConfig, isSending]);
   const isDisabled = useMemo(
-    () =>
-      isUndefined(id) || isUndefined(numberOfVotes) || isError || isLoading || period !== Period.Evidence || isDrawn,
-    [id, numberOfVotes, isError, isLoading, period, isDrawn]
+    () => isUndefined(id) || isUndefined(numberOfVotes) || isError || isLoading || !canDraw,
+    [id, numberOfVotes, isError, isLoading, canDraw]
   );
   const handleClick = () => {
     if (!drawConfig || !publicClient) return;