Merge pull request #3686 from canonical/fix-new-version-unused-restart

Fix `restart` not showing new version
canonical · Sep 27, 2024 · 86053b9 · 86053b9
1 parent 1014151
commit 86053b9
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 27 deletions.
diff --git a/src/client/cli/cmd/restart.cpp b/src/client/cli/cmd/restart.cpp
@@ -39,9 +39,14 @@ mp::ReturnCode cmd::Restart::run(mp::ArgParser* parser)
     if (ret != ParseCode::Ok)
         return parser->returnCodeFrom(ret);
 
-    auto on_success = [](mp::RestartReply& reply) { return ReturnCode::Ok; };
-
     AnimatedSpinner spinner{cout};
+    auto on_success = [this, &spinner](mp::RestartReply& reply) {
+        spinner.stop();
+        if (term->is_live() && update_available(reply.update_info()))
+            cout << update_notice(reply.update_info());
+        return ReturnCode::Ok;
+    };
+
     auto on_failure = [this, &spinner](grpc::Status& status) {
         spinner.stop();
         return standard_failure_handler_for(name(), cerr, status);

diff --git a/src/daemon/daemon.cpp b/src/daemon/daemon.cpp
@@ -3405,25 +3405,19 @@ mp::Daemon::async_wait_for_ready_all(grpc::ServerReaderWriterInterface<Reply, Re
         if (auto error = future.result(); !error.empty())
             add_fmt_to(errors, error);
 
-    if (server && std::is_same<Reply, StartReply>::value)
+    if constexpr (std::is_same_v<Reply, StartReply> || std::is_same_v<Reply, RestartReply>)
     {
-        bool write_reply{false};
-        Reply reply;
-
-        if (config->update_prompt->is_time_to_show())
+        if (server)
         {
-            config->update_prompt->populate(reply.mutable_update_info());
-            write_reply = true;
-        }
+            Reply reply;
 
-        if (warnings.size() > 0)
-        {
-            reply.set_log_line(fmt::to_string(warnings));
-            write_reply = true;
-        }
+            config->update_prompt->populate_if_time_to_show(reply.mutable_update_info());
+
+            if (warnings.size() > 0)
+            {
+                reply.set_log_line(fmt::to_string(warnings));
+            }
 
-        if (write_reply)
-        {
             server->Write(reply);
         }
     }

diff --git a/tests/test_daemon_start.cpp b/tests/test_daemon_start.cpp
@@ -77,8 +77,10 @@ TEST_F(TestDaemonStart, successfulStartOkStatus)
     mp::StartRequest request;
     request.mutable_instance_names()->add_instance_name(mock_instance_name);
 
-    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request,
-                                   StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>>{});
+    StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> mock_server{};
+    EXPECT_CALL(mock_server, Write(_, _)).Times(1);
+
+    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request, std::move(mock_server));
 
     EXPECT_TRUE(status.ok());
 }
@@ -113,8 +115,7 @@ TEST_F(TestDaemonStart, exitlessSshProcessExceptionDoesNotShowMessage)
     request.mutable_instance_names()->add_instance_name(mock_instance_name);
 
     StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> server;
-
-    EXPECT_CALL(server, Write(_, _)).Times(0);
+    EXPECT_CALL(server, Write(_, _)).Times(1);
 
     auto status = call_daemon_slot(daemon, &mp::Daemon::start, request, std::move(server));
 
@@ -143,8 +144,10 @@ TEST_F(TestDaemonStart, unknownStateDoesNotStart)
     mp::StartRequest request;
     request.mutable_instance_names()->add_instance_name(mock_instance_name);
 
-    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request,
-                                   StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>>{});
+    StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> mock_server;
+    EXPECT_CALL(mock_server, Write(_, _)).Times(1);
+
+    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request, std::move(mock_server));
 
     EXPECT_FALSE(status.ok());
 }
@@ -170,8 +173,10 @@ TEST_F(TestDaemonStart, suspendingStateDoesNotStartHasError)
     mp::StartRequest request;
     request.mutable_instance_names()->add_instance_name(mock_instance_name);
 
-    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request,
-                                   StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>>{});
+    StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> mock_server;
+    EXPECT_CALL(mock_server, Write(_, _)).Times(1);
+
+    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request, std::move(mock_server));
 
     EXPECT_FALSE(status.ok());
 
@@ -208,8 +213,10 @@ TEST_F(TestDaemonStart, definedMountsInitializedDuringStart)
     mp::StartRequest request;
     request.mutable_instance_names()->add_instance_name(mock_instance_name);
 
-    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request,
-                                   StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>>{});
+    StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> mock_server;
+    EXPECT_CALL(mock_server, Write(_, _)).Times(1);
+
+    auto status = call_daemon_slot(daemon, &mp::Daemon::start, request, std::move(mock_server));
 
     EXPECT_TRUE(status.ok());
 }
@@ -238,6 +245,8 @@ TEST_F(TestDaemonStart, removingMountOnFailedStart)
 
     auto log = fmt::format("Removing mount \"{}\" from '{}': {}\n", fake_target_path, mock_instance_name, error);
     StrictMock<mpt::MockServerReaderWriter<mp::StartReply, mp::StartRequest>> server;
+
+    EXPECT_CALL(server, Write(_, _));
     EXPECT_CALL(server, Write(Property(&mp::StartReply::log_line, Eq(log)), _));
 
     config_builder.data_directory = temp_dir->path();