Skip to content

Commit

Permalink
feat: add MultiProgress to IndicatifFeedbackReceiver
Browse files Browse the repository at this point in the history
This field is set when `MithrilEventCardanoDatabase::Started` is triggered

Co-authored-by: DJO <[email protected]>
  • Loading branch information
dlachaume and Alenar committed Feb 19, 2025
1 parent 3d209db commit 3ff5aa4
Showing 1 changed file with 174 additions and 5 deletions.
179 changes: 174 additions & 5 deletions mithril-client-cli/src/utils/feedback_receiver.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use async_trait::async_trait;
use indicatif::{ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle};
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle};
use slog::Logger;
use std::fmt::Write;
use tokio::sync::RwLock;
Expand All @@ -13,6 +13,7 @@ use mithril_client::feedback::{FeedbackReceiver, MithrilEvent, MithrilEventCarda
pub struct IndicatifFeedbackReceiver {
download_progress_reporter: RwLock<Option<DownloadProgressReporter>>,
certificate_validation_pb: RwLock<Option<ProgressBar>>,
cardano_database_multi_pb: RwLock<Option<MultiProgress>>,
output_type: ProgressOutputType,
logger: Logger,
}
Expand All @@ -23,6 +24,7 @@ impl IndicatifFeedbackReceiver {
Self {
download_progress_reporter: RwLock::new(None),
certificate_validation_pb: RwLock::new(None),
cardano_database_multi_pb: RwLock::new(None),
output_type,
logger,
}
Expand Down Expand Up @@ -72,15 +74,25 @@ impl FeedbackReceiver for IndicatifFeedbackReceiver {
*download_progress_reporter = None;
}
MithrilEvent::CardanoDatabase(cardano_database_event) => match cardano_database_event {
MithrilEventCardanoDatabase::Started { download_id: _ } => {
println!("MithrilEventCardanoDatabase::Started")
MithrilEventCardanoDatabase::Started {
download_id: _,
total_immutable_files: _,
include_ancillary: _,
} => {
let multi_pb = MultiProgress::new();
let mut cardano_database_multi_pb =
self.cardano_database_multi_pb.write().await;
*cardano_database_multi_pb = Some(multi_pb);
}
MithrilEventCardanoDatabase::Completed { download_id: _ } => {
println!("MithrilEventCardanoDatabase::Completed")
let mut cardano_database_multi_pb =
self.cardano_database_multi_pb.write().await;
*cardano_database_multi_pb = None;
}
MithrilEventCardanoDatabase::ImmutableDownloadStarted {
immutable_file_number: _,
download_id: _,
size: _,
} => {
println!("MithrilEventCardanoDatabase::ImmutableDownloadStarted")
}
Expand All @@ -98,7 +110,10 @@ impl FeedbackReceiver for IndicatifFeedbackReceiver {
} => {
println!("MithrilEventCardanoDatabase::ImmutableDownloadCompleted")
}
MithrilEventCardanoDatabase::AncillaryDownloadStarted { download_id: _ } => {
MithrilEventCardanoDatabase::AncillaryDownloadStarted {
download_id: _,
size: _,
} => {
println!("MithrilEventCardanoDatabase::AncillaryDownloadStarted")
}
MithrilEventCardanoDatabase::AncillaryDownloadProgress {
Expand Down Expand Up @@ -168,3 +183,157 @@ impl FeedbackReceiver for IndicatifFeedbackReceiver {
}
}
}

#[cfg(test)]
mod tests {
use slog::o;
use std::sync::Arc;

use mithril_client::feedback::FeedbackSender;

use super::*;

const DOWNLOAD_ID: &str = "id";

macro_rules! send_event {
(cardano_db, dl_started => $sender:expr, $total_immutable:expr, $include_ancillary:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::Started {
download_id: DOWNLOAD_ID.to_string(),
total_immutable_files: $total_immutable,
include_ancillary: $include_ancillary,
},
))
.await;
};
(cardano_db, dl_completed => $sender:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::Completed {
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, immutable_dl, started => $sender:expr, immutable:$immutable_file_number:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::ImmutableDownloadStarted {
immutable_file_number: $immutable_file_number,
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, immutable_dl, progress => $sender:expr, immutable:$immutable_file_number:expr, bytes:$downloaded_bytes:expr, size:$size:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::ImmutableDownloadProgress {
immutable_file_number: $immutable_file_number,
download_id: DOWNLOAD_ID.to_string(),
downloaded_bytes: $downloaded_bytes,
size: $size,
},
))
.await;
};
(cardano_db, immutable_dl, completed => $sender:expr, immutable:$immutable_file_number:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::ImmutableDownloadCompleted {
immutable_file_number: $immutable_file_number,
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, ancillary_dl, started => $sender:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::AncillaryDownloadStarted {
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, ancillary_dl, progress => $sender:expr, bytes:$downloaded_bytes:expr, size:$size:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::AncillaryDownloadProgress {
download_id: DOWNLOAD_ID.to_string(),
downloaded_bytes: $downloaded_bytes,
size: $size,
},
))
.await;
};
(cardano_db, ancillary_dl, completed => $sender:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::AncillaryDownloadCompleted {
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, digests_dl, started => $sender:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::DigestDownloadStarted {
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
(cardano_db, digests_dl, progress => $sender:expr, bytes:$downloaded_bytes:expr, size:$size:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::DigestDownloadProgress {
download_id: DOWNLOAD_ID.to_string(),
downloaded_bytes: $downloaded_bytes,
size: $size,
},
))
.await;
};
(cardano_db, digests_dl, completed => $sender:expr) => {
$sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::DigestDownloadCompleted {
download_id: DOWNLOAD_ID.to_string(),
},
))
.await;
};
}

#[tokio::test]
async fn starting_a_cardano_database_should_add_multi_progress_bar() {
let receiver = Arc::new(IndicatifFeedbackReceiver::new(
ProgressOutputType::Hidden,
slog::Logger::root(slog::Discard, o!()),
));
let sender = FeedbackSender::new(&[receiver.clone()]);

send_event!(cardano_db, dl_started => sender, 99, false);
send_event!(cardano_db, immutable_dl, progress => sender, immutable:2, bytes:12, size:43);

assert!(receiver.cardano_database_multi_pb.read().await.is_some());
}

#[tokio::test]
async fn completing_a_cardano_database_when_started_event_before_should_remove_multi_progress_bar(
) {
let receiver = Arc::new(IndicatifFeedbackReceiver::new(
ProgressOutputType::Hidden,
slog::Logger::root(slog::Discard, o!()),
));
let sender = FeedbackSender::new(&[receiver.clone()]);

send_event!(cardano_db, dl_started => sender, 99, false);
send_event!(cardano_db, dl_completed => sender);

assert!(receiver.cardano_database_multi_pb.read().await.is_none());
}
}

0 comments on commit 3ff5aa4

Please sign in to comment.