diff --git a/service/history/api/queryworkflow/api.go b/service/history/api/queryworkflow/api.go
index 5aa100bea95..7e75e98e04f 100644
--- a/service/history/api/queryworkflow/api.go
+++ b/service/history/api/queryworkflow/api.go
@@ -115,7 +115,7 @@ func Invoke(
 		return nil, consts.ErrWorkflowTaskNotScheduled
 	}
 
-	if mutableState.GetExecutionInfo().WorkflowTaskAttempt > 1 {
+	if mutableState.IsTransientWorkflowTask() {
 		// while workflow task is failing, the query to that workflow will also fail. Failing fast here to prevent wasting
 		// resources to load history for a query that will fail.
 		return nil, serviceerror.NewFailedPrecondition("Cannot query workflow due to Workflow Task in failed state.")
diff --git a/tests/query_workflow_test.go b/tests/query_workflow_test.go
index 7e3b42dc9c8..5ea849c9f1b 100644
--- a/tests/query_workflow_test.go
+++ b/tests/query_workflow_test.go
@@ -267,7 +267,7 @@ func (s *clientIntegrationSuite) TestQueryWorkflow_QueryFailedWorkflowTask() {
 		})
 
 		if err != nil {
-			s.Logger.Fatal("SetQueryHandler failed: " + err.Error())
+			s.T().Fatalf("SetQueryHandler failed: %s", err.Error())
 		}
 		// force workflow task to fail
 		panic("Workflow failed")
@@ -292,7 +292,8 @@ func (s *clientIntegrationSuite) TestQueryWorkflow_QueryFailedWorkflowTask() {
 	s.True(workflowRun.GetRunID() != "")
 
 	// wait for workflow to fail
-	time.Sleep(time.Second * 5)
+	time.Sleep(time.Second * 2)
 	_, err = s.sdkClient.QueryWorkflow(ctx, id, "", "test")
+	s.Error(err)
 	s.IsType(&serviceerror.FailedPrecondition{}, err)
 }