[test] preempt stages for early stop

Author: captain-yoshi

core/test/test_container.cpp

@@ -674,21 +674,20 @@ TEST(Task, timeout) {
 }
 
 // https://github.com/moveit/moveit_task_constructor/pull/597
+// https://github.com/moveit/moveit_task_constructor/pull/598
 // start planning in another thread, then preempt it in this thread
-TEST(Task, preempt) {
+TEST_F(TaskTestBase, preempt) {
 	moveit::core::MoveItErrorCode ec;
 	resetMockupIds();
 
-	Task t;
-	t.setRobotModel(getModel());
-
 	auto timeout = std::chrono::milliseconds(10);
-	t.add(std::make_unique<GeneratorMockup>(PredefinedCosts::constant(0.0)));
-	t.add(std::make_unique<TimedForwardMockup>(timeout));
+	auto gen1 = add(t, new GeneratorMockup(PredefinedCosts::constant(0.0)));
+	auto fwd1 = add(t, new TimedForwardMockup(timeout));
+	auto fwd2 = add(t, new TimedForwardMockup(timeout));
 
 	// preempt before preempt_request_ is reset in plan()
 	{
-		std::thread thread{ [&ec, &t, timeout] {
+		std::thread thread{ [&ec, this, timeout] {
 			std::this_thread::sleep_for(timeout);
 			ec = t.plan(1);
 		} };
@@ -698,5 +697,30 @@ TEST(Task, preempt) {
 
 	EXPECT_EQ(ec, moveit::core::MoveItErrorCode::PREEMPTED);
 	EXPECT_EQ(t.solutions().size(), 0u);
+	EXPECT_EQ(gen1->runs_, 0u);
+	EXPECT_EQ(fwd1->runs_, 0u);
+	EXPECT_EQ(fwd2->runs_, 0u);
+	EXPECT_TRUE(t.plan(1));  // make sure the preempt request has been resetted on the previous call to plan()
+
+	// prempt in between stage computation
+	// Failing stage(s):
+	// 0  - ←   0 →   -  0 / task pipeline
+	//    1  - ←   1 →   -  0 / GEN1
+	//    -  0 →   1 →   -  1 / FWD1
+	//    -  1 →   0 →   -  0 / FWD2
+	t.reset();
+	resetMockupIds();
+	{
+		std::thread thread{ [&ec, this] { ec = t.plan(1); } };
+		std::this_thread::sleep_for(timeout);
+		t.preempt();
+		thread.join();
+	}
+
+	EXPECT_EQ(ec, moveit::core::MoveItErrorCode::PREEMPTED);
+	EXPECT_EQ(t.solutions().size(), 0u);
+	EXPECT_EQ(gen1->runs_, 2u);
+	EXPECT_EQ(fwd1->runs_, 2u);
+	EXPECT_EQ(fwd2->runs_, 1u);
 	EXPECT_TRUE(t.plan(1));  // make sure the preempt request has been resetted on the previous call to plan()
 }