Merge e6d68be into d218dd6

UpdateHub · Oct 19, 2020 · b8d5647 · b8d5647
2 parents d218dd6 + e6d68be
commit b8d5647
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/updatehub-sdk/src/listener.rs b/updatehub-sdk/src/listener.rs
@@ -19,6 +19,7 @@ type CallbackFn = dyn Fn(Handler) -> Pin<Box<dyn Future<Output = Result<()>>>>;
 
 #[derive(Default)]
 pub struct StateChangeListener {
+    probe_callbacks: Vec<Box<CallbackFn>>,
     download_callbacks: Vec<Box<CallbackFn>>,
     install_callbacks: Vec<Box<CallbackFn>>,
     reboot_callbacks: Vec<Box<CallbackFn>>,
@@ -27,6 +28,7 @@ pub struct StateChangeListener {
 
 #[derive(Debug)]
 pub enum State {
+    Probe,
     Download,
     Install,
     Reboot,
@@ -63,6 +65,7 @@ impl StateChangeListener {
         Fut: Future<Output = Result<()>> + 'static,
     {
         match state {
+            State::Probe => self.probe_callbacks.push(Box::new(move |d| Box::pin(f(d)))),
             State::Download => self.download_callbacks.push(Box::new(move |d| Box::pin(f(d)))),
             State::Install => self.install_callbacks.push(Box::new(move |d| Box::pin(f(d)))),
             State::Reboot => self.reboot_callbacks.push(Box::new(move |d| Box::pin(f(d)))),
@@ -100,6 +103,7 @@ impl StateChangeListener {
 
     async fn emit(&self, stream: UnixStream, input: &str) -> Result<()> {
         let callbacks = match input {
+            "probe" => &self.probe_callbacks,
             "download" => &self.download_callbacks,
             "install" => &self.install_callbacks,
             "reboot" => &self.reboot_callbacks,

diff --git a/updatehub/src/states/mod.rs b/updatehub/src/states/mod.rs
@@ -250,7 +250,7 @@ impl State {
             State::Park(s) => s.handle(context).await,
             State::EntryPoint(s) => s.handle(context).await,
             State::Poll(s) => s.handle(context).await,
-            State::Probe(s) => s.handle(context).await,
+            State::Probe(s) => s.handle_with_callback(context).await,
             State::Validation(s) => s.handle(context).await,
             State::DirectDownload(s) => s.handle(context).await,
             State::PrepareLocalInstall(s) => s.handle(context).await,

diff --git a/updatehub/src/states/probe.rs b/updatehub/src/states/probe.rs
@@ -4,7 +4,7 @@
 
 use super::{
     machine::{self, Context},
-    EntryPoint, Result, State, StateChangeImpl, Validation,
+    CallbackReporter, EntryPoint, Result, State, StateChangeImpl, Validation,
 };
 use chrono::{Duration, Utc};
 use cloud::api::ProbeResponse;
@@ -13,6 +13,8 @@ use slog_scope::{error, info};
 #[derive(Debug)]
 pub(super) struct Probe;
 
+impl CallbackReporter for Probe {}
+
 /// Implements the state change for State<Probe>.
 #[async_trait::async_trait]
 impl StateChangeImpl for Probe {

diff --git a/updatehub/tests/successful_integration_test.rs b/updatehub/tests/successful_integration_test.rs
@@ -329,7 +329,7 @@ fn correct_config_update_polling() {
 #[test]
 fn correct_config_statechange_callback() {
     let state_change_script = r#"#! /bin/bash
-echo "cancel"
+[ "$1" = "download" ] && echo "cancel" || echo
 "#;
 
     let (mut session, setup) = Settings::default()
@@ -391,7 +391,7 @@ echo "cancel"
 #[test]
 fn correct_config_error_state_callback() {
     let state_change_script = r#"#! /bin/bash
-echo "cancel"
+[ "$1" = "error" ] && echo "cancel" || echo
 "#;
 
     let (mut session, setup) = Settings::default()