Merge pull request #511 from Dronecode/fix-mission-progress-tests

backend: fix SubscribeMissionProgress unit tests
mavlink · Aug 29, 2018 · e17c3be · e17c3be
2 parents 83a7399 + bac618a
commit e17c3be
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 23 deletions.
diff --git a/backend/src/plugins/mission/mission_service_impl.h b/backend/src/plugins/mission/mission_service_impl.h
@@ -278,23 +278,24 @@ class MissionServiceImpl final : public dronecode_sdk::rpc::mission::MissionServ
         const dronecode_sdk::rpc::mission::SubscribeMissionProgressRequest * /* request */,
         grpc::ServerWriter<rpc::mission::MissionProgressResponse> *writer) override
     {
-        std::promise<void> mission_finished_promise;
-        auto mission_finished_future = mission_finished_promise.get_future();
-
-        _mission.subscribe_progress([&writer](int current, int total) {
-            dronecode_sdk::rpc::mission::MissionProgressResponse rpc_mission_progress_response;
-            rpc_mission_progress_response.set_current_item_index(current);
-            rpc_mission_progress_response.set_mission_count(total);
-            writer->Write(rpc_mission_progress_response);
-
-            // FIXME: It is possible that we get this progress update multiple times. If this
-            //        happens, we fulfill the promise more than once which leads to a crash.
-            // if (current == total - 1) {
-            //    mission_finished_promise.set_value();
-            //}
-        });
-
-        mission_finished_future.wait();
+        std::promise<void> stream_closed_promise;
+        auto stream_closed_future = stream_closed_promise.get_future();
+
+        bool is_finished = false;
+
+        _mission.subscribe_progress(
+            [&writer, &stream_closed_promise, &is_finished](int current, int total) {
+                dronecode_sdk::rpc::mission::MissionProgressResponse rpc_mission_progress_response;
+                rpc_mission_progress_response.set_current_item_index(current);
+                rpc_mission_progress_response.set_mission_count(total);
+
+                if (!writer->Write(rpc_mission_progress_response) && !is_finished) {
+                    is_finished = true;
+                    stream_closed_promise.set_value();
+                }
+            });
+
+        stream_closed_future.wait();
         return grpc::Status::OK;
     }
 

diff --git a/backend/test/mission_service_impl_test.cpp b/backend/test/mission_service_impl_test.cpp
@@ -579,7 +579,8 @@ class MissionServiceImplProgressTest : public MissionServiceImplTestBase {
     }
 
     std::future<void>
-    subscribeMissionProgressAsync(std::vector<std::pair<int, int>> &progress_events) const;
+    subscribeMissionProgressAsync(std::vector<std::pair<int, int>> &progress_events,
+                                  std::shared_ptr<grpc::ClientContext> context) const;
 
     std::unique_ptr<grpc::Server> _server;
     std::unique_ptr<MissionService::Stub> _stub;
@@ -588,23 +589,26 @@ class MissionServiceImplProgressTest : public MissionServiceImplTestBase {
 TEST_F(MissionServiceImplProgressTest, registersToMissionProgress)
 {
     dc::Mission::progress_callback_t progress_callback;
+    auto context = std::make_shared<grpc::ClientContext>();
     EXPECT_CALL(_mission, subscribe_progress(_))
         .WillOnce(SaveResult(&progress_callback, &_callback_saved_promise));
     std::vector<std::pair<int, int>> progress_events;
 
-    auto progress_events_future = subscribeMissionProgressAsync(progress_events);
+    auto progress_events_future = subscribeMissionProgressAsync(progress_events, context);
     _callback_saved_future.wait();
     progress_callback(0, 1);
+    context->TryCancel();
+    progress_callback(0, 0); // TryCancel() requires one more event to trigger...
     progress_events_future.wait();
 }
 
 std::future<void> MissionServiceImplProgressTest::subscribeMissionProgressAsync(
-    std::vector<std::pair<int, int>> &progress_events) const
+    std::vector<std::pair<int, int>> &progress_events,
+    std::shared_ptr<grpc::ClientContext> context) const
 {
     return std::async(std::launch::async, [&]() {
-        grpc::ClientContext context;
         dronecode_sdk::rpc::mission::SubscribeMissionProgressRequest request;
-        auto response_reader = _stub->SubscribeMissionProgress(&context, request);
+        auto response_reader = _stub->SubscribeMissionProgress(context.get(), request);
 
         dronecode_sdk::rpc::mission::MissionProgressResponse response;
         while (response_reader->Read(&response)) {
@@ -631,12 +635,15 @@ TEST_F(MissionServiceImplProgressTest, SendsMultipleMissionProgressEvents)
     }
     std::vector<std::pair<int, int>> received_progress_events;
 
-    auto progress_events_future = subscribeMissionProgressAsync(received_progress_events);
+    auto context = std::make_shared<grpc::ClientContext>();
+    auto progress_events_future = subscribeMissionProgressAsync(received_progress_events, context);
     _callback_saved_future.wait();
 
     for (const auto progress_event : expected_progress_events) {
         progress_callback(progress_event.first, progress_event.second);
     }
+    context->TryCancel();
+    progress_callback(0, 0); // TryCancel() requires one more event to trigger...
     progress_events_future.wait();
 
     ASSERT_EQ(expected_mission_count, received_progress_events.size());