Merge pull request #970 from othrayte/improve-e2e-test-reliability

website: Improve reliability of e2e tests

website/src/components/Loading/LoadingScreen.jsx

@@ -1,4 +1,4 @@
-import { Box, Center, Progress, Text, useColorModeValue } from "@chakra-ui/react";
+import { Box, Center, Progress, Text } from "@chakra-ui/react";
 
 export const LoadingScreen = ({ text = "Loading..." } = {}) => {
   return (

website/src/components/Survey/TaskControls.tsx

@@ -1,4 +1,4 @@
-import { Box, Flex, IconButton, Tooltip, useColorModeValue } from "@chakra-ui/react";
+import { Box, Flex, IconButton, Progress, Tooltip, useColorModeValue } from "@chakra-ui/react";
 import { Edit2 } from "lucide-react";
 import { SkipButton } from "src/components/Buttons/Skip";
 import { SubmitButton } from "src/components/Buttons/Submit";
@@ -9,55 +9,58 @@ import { BaseTask } from "src/types/Task";
 export interface TaskControlsProps {
   task: BaseTask;
   taskStatus: TaskStatus;
+  isLoading: boolean;
   onEdit: () => void;
   onReview: () => void;
   onSubmit: () => void;
   onSkip: (reason: string) => void;
 }
 
-export const TaskControls = ({ task, taskStatus, onEdit, onReview, onSubmit, onSkip }: TaskControlsProps) => {
+export const TaskControls = ({
+  task,
+  taskStatus,
+  isLoading,
+  onEdit,
+  onReview,
+  onSubmit,
+  onSkip,
+}: TaskControlsProps) => {
   const backgroundColor = useColorModeValue("white", "gray.800");
 
   return (
-    <Box
-      width="full"
-      bg={backgroundColor}
-      borderRadius="xl"
-      p="6"
-      display="flex"
-      flexDirection={["column", "row"]}
-      shadow="base"
-      gap="4"
-    >
-      <TaskInfo id={task.id} output="Submit your answer" />
-      <Flex width={["full", "fit-content"]} justify="center" ml="auto" gap={2}>
-        {taskStatus.mode === "EDIT" ? (
-          <>
-            <SkipButton onSkip={onSkip} />
-            <SubmitButton
-              colorScheme="blue"
-              data-cy="review"
-              isDisabled={taskStatus.replyValidity === "INVALID"}
-              onClick={onReview}
-            >
-              Review
-            </SubmitButton>
-          </>
-        ) : (
-          <>
-            <Tooltip label="Edit">
-              <IconButton size="lg" data-cy="edit" aria-label="edit" onClick={onEdit} icon={<Edit2 size="1em" />} />
-            </Tooltip>
-            <SubmitButton
-              colorScheme="green"
-              data-cy="submit"
-              isDisabled={taskStatus.mode === "SUBMITTED"}
-              onClick={onSubmit}
-            >
-              Submit
-            </SubmitButton>
-          </>
-        )}
+    <Box width="full" bg={backgroundColor} borderRadius="xl" shadow="base">
+      {isLoading && <Progress size="sm" isIndeterminate />}
+      <Flex p="6" gap="4" direction={["column", "row"]}>
+        <TaskInfo id={task.id} output="Submit your answer" />
+        <Flex width={["full", "fit-content"]} justify="center" ml="auto" gap={2}>
+          {taskStatus.mode === "EDIT" ? (
+            <>
+              <SkipButton onSkip={onSkip} />
+              <SubmitButton
+                colorScheme="blue"
+                data-cy="review"
+                isDisabled={taskStatus.replyValidity === "INVALID"}
+                onClick={onReview}
+              >
+                Review
+              </SubmitButton>
+            </>
+          ) : (
+            <>
+              <Tooltip label="Edit">
+                <IconButton size="lg" data-cy="edit" aria-label="edit" onClick={onEdit} icon={<Edit2 size="1em" />} />
+              </Tooltip>
+              <SubmitButton
+                colorScheme="green"
+                data-cy="submit"
+                isDisabled={taskStatus.mode === "SUBMITTED"}
+                onClick={onSubmit}
+              >
+                Submit
+              </SubmitButton>
+            </>
+          )}
+        </Flex>
       </Flex>
     </Box>
   );

website/src/components/TaskPage/TaskPage.tsx

@@ -4,9 +4,10 @@ import { TaskEmptyState } from "src/components/EmptyState";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
 import { TaskInfos } from "src/components/Tasks/TaskTypes";
-import { ERROR_CODES, taskApiHooks } from "src/lib/constants";
+import { taskApiHooks } from "src/lib/constants";
 import { getTypeSafei18nKey } from "src/lib/i18n";
 import { TaskType } from "src/types/Task";
+import { KnownTaskType } from "src/types/Tasks";
 
 type TaskPageProps = {
   type: TaskType;
@@ -15,26 +16,38 @@ type TaskPageProps = {
 export const TaskPage = ({ type }: TaskPageProps) => {
   const { t } = useTranslation(["tasks", "common"]);
   const taskApiHook = taskApiHooks[type];
-  const { tasks, isLoading, reset, trigger, error } = taskApiHook(type);
+  const { response, isLoading, completeTask, skipTask } = taskApiHook(type);
   const taskInfo = TaskInfos.find((taskType) => taskType.type === type);
 
-  if (isLoading) {
-    return <LoadingScreen text={t("common:loading")} />;
+  let body;
+  switch (response.taskAvailability) {
+    case "AWAITING_INITIAL":
+      body = <LoadingScreen text={t("common:loading")} />;
+      break;
+    case "NONE_AVAILABLE":
+      body = <TaskEmptyState />;
+      break;
+    case "AVAILABLE":
+      body = (
+        <Task
+          key={response.task.id}
+          frontendId={response.id}
+          task={response.task as KnownTaskType}
+          isLoading={isLoading}
+          completeTask={completeTask}
+          skipTask={skipTask}
+        />
+      );
+      break;
   }
 
-  if (tasks.length === 0 || error?.errorCode === ERROR_CODES.TASK_REQUESTED_TYPE_NOT_AVAILABLE) {
-    return <TaskEmptyState />;
-  }
-
-  const task = tasks[0];
-
   return (
     <>
       <Head>
         <title>{t(getTypeSafei18nKey(`${taskInfo.id}.label`))}</title>
         <meta name="description" content={t(getTypeSafei18nKey(`${taskInfo.id}.desc`))} />
       </Head>
-      <Task key={task.task.id} frontendId={task.id} task={task.task} trigger={trigger} mutate={reset} />
+      {body}
     </>
   );
 };

website/src/components/Tasks/CreateTask.tsx

@@ -8,7 +8,7 @@ import { TaskSurveyProps } from "src/components/Tasks/Task";
 import { TaskHeader } from "src/components/Tasks/TaskHeader";
 import { getTypeSafei18nKey } from "src/lib/i18n";
 import { TaskType } from "src/types/Task";
-import { CreateAssistantReplyTask, CreateInitialPromptTask, CreatePrompterReplyTask } from "src/types/Tasks";
+import { CreateTaskType } from "src/types/Tasks";
 
 export const CreateTask = ({
   task,
@@ -17,7 +17,7 @@ export const CreateTask = ({
   isDisabled,
   onReplyChanged,
   onValidityChanged,
-}: TaskSurveyProps<CreateInitialPromptTask | CreateAssistantReplyTask | CreatePrompterReplyTask, { text: string }>) => {
+}: TaskSurveyProps<CreateTaskType, { text: string }>) => {
   const { t, i18n } = useTranslation(["tasks", "common"]);
   const cardColor = useColorModeValue("gray.50", "gray.800");
   const titleColor = useColorModeValue("gray.800", "gray.300");

website/src/components/Tasks/EvaluateTask.tsx

@@ -6,7 +6,7 @@ import { SurveyCard } from "src/components/Survey/SurveyCard";
 import { TaskSurveyProps } from "src/components/Tasks/Task";
 import { TaskHeader } from "src/components/Tasks/TaskHeader";
 import { TaskType } from "src/types/Task";
-import { RankAssistantRepliesTask, RankInitialPromptsTask, RankPrompterRepliesTask } from "src/types/Tasks";
+import { RankTaskType } from "src/types/Tasks";
 
 export const EvaluateTask = ({
   task,
@@ -15,10 +15,7 @@ export const EvaluateTask = ({
   isDisabled,
   onReplyChanged,
   onValidityChanged,
-}: TaskSurveyProps<
-  RankInitialPromptsTask | RankAssistantRepliesTask | RankPrompterRepliesTask,
-  { ranking: number[] }
->) => {
+}: TaskSurveyProps<RankTaskType, { ranking: number[] }>) => {
   const cardColor = useColorModeValue("gray.50", "gray.800");
   const [ranking, setRanking] = useState<number[]>(null);
 

website/src/components/Tasks/Task/Task.stories.tsx

@@ -7,8 +7,10 @@ export default {
   component: Task,
 };
 
-const Template = ({ frontendId, task, trigger, mutate }) => {
-  return <Task frontendId={frontendId} task={task} trigger={trigger} mutate={mutate} />;
+const Template = ({ frontendId, task, isLoading, completeTask, skipTask }) => {
+  return (
+    <Task frontendId={frontendId} task={task} isLoading={isLoading} completeTask={completeTask} skipTask={skipTask} />
+  );
 };
 
 export const Default = Template.bind({});
@@ -23,10 +25,11 @@ Default.args = {
     type: "label_prompter_reply",
     valid_labels: ["spam", "fails_task"],
   },
-  trigger: (id, update_type, content) => {
+  isLoading: false,
+  completeTask: (id, update_type, content) => {
     console.log(content);
   },
-  mutate: () => {
-    console.log("mutate");
+  skipTask: () => {
+    console.log("skip");
   },
 };

website/src/components/Tasks/Task/Task.tsx

@@ -10,6 +10,7 @@ import { UnchangedWarning } from "src/components/Tasks/UnchangedWarning";
 import { post } from "src/lib/api";
 import { getTypeSafei18nKey } from "src/lib/i18n";
 import { BaseTask, TaskContent, TaskReplyValidity } from "src/types/Task";
+import { CreateTaskType, KnownTaskType, LabelTaskType, RankTaskType } from "src/types/Tasks";
 import useSWRMutation from "swr/mutation";
 
 interface EditMode {
@@ -62,7 +63,15 @@ export interface TaskSurveyProps<TaskType extends BaseTask, T> {
   onValidityChanged: (validity: TaskReplyValidity) => void;
 }
 
-export const Task = ({ frontendId, task, trigger, mutate }) => {
+interface TaskProps {
+  frontendId: string;
+  task: KnownTaskType;
+  isLoading: boolean;
+  completeTask: (TaskContent) => void;
+  skipTask: () => void;
+}
+
+export const Task = ({ frontendId, task, isLoading, completeTask, skipTask }: TaskProps) => {
   const { t } = useTranslation("tasks");
   const [taskStatus, taskEvent] = useReducer(
     (
@@ -117,14 +126,13 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
 
   const rootEl = useRef<HTMLDivElement>(null);
 
-  const taskType = useMemo(
-    () => TaskInfos.find((taskType) => taskType.type === task.type && taskType.mode === task.mode),
-    [task.type, task.mode]
-  );
+  const taskType = useMemo(() => {
+    return TaskInfos.find((taskType) => taskType.type === task.type);
+  }, [task.type]);
 
   const { trigger: sendRejection } = useSWRMutation("/api/reject_task", post, {
     onSuccess: async () => {
-      mutate();
+      skipTask();
     },
   });
 
@@ -144,7 +152,7 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
 
   const submitResponse = () => {
     if (taskStatus.mode === "REVIEW") {
-      trigger({
+      completeTask({
         id: frontendId,
         update_type: taskType.update_type,
         content: replyContent.current,
@@ -159,7 +167,7 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
       case TaskCategory.Create:
         return (
           <CreateTask
-            task={task}
+            task={task as CreateTaskType}
             taskType={taskType}
             isEditable={taskStatus.mode === "EDIT"}
             isDisabled={taskStatus.mode === "SUBMITTED"}
@@ -170,7 +178,7 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
       case TaskCategory.Evaluate:
         return (
           <EvaluateTask
-            task={task}
+            task={task as RankTaskType}
             taskType={taskType}
             isEditable={taskStatus.mode === "EDIT"}
             isDisabled={taskStatus.mode === "SUBMITTED"}
@@ -181,7 +189,7 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
       case TaskCategory.Label:
         return (
           <LabelTask
-            task={task}
+            task={task as LabelTaskType}
             taskType={taskType}
             isEditable={taskStatus.mode === "EDIT"}
             isDisabled={taskStatus.mode === "SUBMITTED"}
@@ -198,6 +206,7 @@ export const Task = ({ frontendId, task, trigger, mutate }) => {
       <TaskControls
         task={task}
         taskStatus={taskStatus}
+        isLoading={isLoading}
         onEdit={() => taskEvent({ action: "RETURN_EDIT" })}
         onReview={() => taskEvent({ action: "REVIEW" })}
         onSubmit={submitResponse}