mithril_client_cli/commands/tools/
aggregator_discovery.rs

1use clap::Parser;
2use cli_table::{Cell, Table, print_stdout};
3
4use mithril_client::{
5    AggregatorDiscoveryType, ClientBuilder, MithrilResult, RequiredAggregatorCapabilities,
6    common::{AggregateSignatureType, MithrilNetwork, SignedEntityTypeDiscriminants},
7};
8
9use crate::CommandContext;
10
11/// Clap command to select an aggregator from the available ones with automatic discovery.
12#[derive(Parser, Debug, Clone)]
13pub struct AggregatorDiscoveryCommand {
14    /// Mithril network name
15    network: MithrilNetwork,
16
17    /// Maximum number of entries to retrieve
18    #[clap(long, default_value_t = 1)]
19    max_entries: usize,
20
21    /// Signed entity types to consider for the discovery
22    ///
23    /// If not provided, all signed entity types are considered.
24    #[clap(long, value_parser, num_args = 0.., value_delimiter = ',')]
25    signed_entity_types: Vec<SignedEntityTypeDiscriminants>,
26
27    /// Aggregate signature types to consider for the discovery
28    ///
29    /// If not provided, all aggregate signature types are considered.
30    #[clap(long, value_parser, num_args = 0.., value_delimiter = ',')]
31    aggregate_signature_types: Vec<AggregateSignatureType>,
32}
33
34impl AggregatorDiscoveryCommand {
35    /// Main command execution
36    pub async fn execute(&self, context: CommandContext) -> MithrilResult<()> {
37        let required_capabilities = self.build_required_capabilities();
38        let client_builder =
39            ClientBuilder::new(AggregatorDiscoveryType::Automatic(self.network.clone()))
40                .with_capabilities(required_capabilities);
41        let aggregator_endpoints = client_builder
42            .discover_aggregator(&self.network)?
43            .take(self.max_entries);
44        let lines = aggregator_endpoints.collect::<Vec<_>>();
45
46        if context.is_json_output_enabled() {
47            println!("{}", serde_json::to_string(&lines)?);
48        } else {
49            let lines = lines
50                .into_iter()
51                .map(|endpoint| vec![endpoint.cell()])
52                .collect::<Vec<_>>()
53                .table()
54                .title(vec!["Aggregator Endpoint".cell()]);
55            print_stdout(lines)?;
56        }
57
58        Ok(())
59    }
60
61    fn build_required_capabilities(&self) -> RequiredAggregatorCapabilities {
62        if self.signed_entity_types.is_empty() && self.aggregate_signature_types.is_empty() {
63            return RequiredAggregatorCapabilities::All;
64        }
65
66        let mut required_capabilities = vec![];
67        if !self.signed_entity_types.is_empty() {
68            let mut required_capabilities_signed_entity_types = vec![];
69            for signed_entity_type in &self.signed_entity_types {
70                required_capabilities_signed_entity_types.push(
71                    RequiredAggregatorCapabilities::SignedEntityType(*signed_entity_type),
72                );
73            }
74            required_capabilities.push(RequiredAggregatorCapabilities::And(
75                required_capabilities_signed_entity_types,
76            ));
77        }
78
79        if !self.aggregate_signature_types.is_empty() {
80            let mut required_capabilities_aggregate_signature_types = vec![];
81            for aggregate_signature_type in &self.aggregate_signature_types {
82                required_capabilities_aggregate_signature_types.push(
83                    RequiredAggregatorCapabilities::AggregateSignatureType(
84                        *aggregate_signature_type,
85                    ),
86                );
87            }
88            required_capabilities.push(RequiredAggregatorCapabilities::And(
89                required_capabilities_aggregate_signature_types,
90            ));
91        }
92        if required_capabilities.len() == 1 {
93            required_capabilities.into_iter().next().unwrap()
94        } else {
95            RequiredAggregatorCapabilities::And(required_capabilities)
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use mithril_client::common::SignedEntityTypeDiscriminants;
104
105    #[test]
106    fn test_build_required_capabilities_all() {
107        let command = AggregatorDiscoveryCommand {
108            network: MithrilNetwork::dummy(),
109            max_entries: 1,
110            signed_entity_types: vec![],
111            aggregate_signature_types: vec![],
112        };
113
114        let required_capabilities = command.build_required_capabilities();
115        assert_eq!(required_capabilities, RequiredAggregatorCapabilities::All);
116    }
117
118    #[test]
119    fn test_build_required_capabilities_signed_entity_types() {
120        let command = AggregatorDiscoveryCommand {
121            network: MithrilNetwork::dummy(),
122            max_entries: 1,
123            signed_entity_types: vec![
124                SignedEntityTypeDiscriminants::CardanoTransactions,
125                SignedEntityTypeDiscriminants::CardanoStakeDistribution,
126            ],
127            aggregate_signature_types: vec![],
128        };
129
130        let required_capabilities = command.build_required_capabilities();
131
132        assert_eq!(
133            required_capabilities,
134            RequiredAggregatorCapabilities::And(vec![
135                RequiredAggregatorCapabilities::SignedEntityType(
136                    SignedEntityTypeDiscriminants::CardanoTransactions
137                ),
138                RequiredAggregatorCapabilities::SignedEntityType(
139                    SignedEntityTypeDiscriminants::CardanoStakeDistribution
140                ),
141            ])
142        );
143    }
144
145    #[test]
146    fn test_build_required_capabilities_aggregate_signature_types() {
147        let command = AggregatorDiscoveryCommand {
148            network: MithrilNetwork::dummy(),
149            max_entries: 1,
150            signed_entity_types: vec![],
151            aggregate_signature_types: vec![AggregateSignatureType::Concatenation],
152        };
153        let required_capabilities = command.build_required_capabilities();
154
155        assert_eq!(
156            required_capabilities,
157            RequiredAggregatorCapabilities::And(vec![
158                RequiredAggregatorCapabilities::AggregateSignatureType(
159                    AggregateSignatureType::Concatenation
160                ),
161            ])
162        );
163    }
164
165    #[test]
166    fn test_build_required_capabilities_both() {
167        let command = AggregatorDiscoveryCommand {
168            network: MithrilNetwork::dummy(),
169            max_entries: 1,
170            signed_entity_types: vec![
171                SignedEntityTypeDiscriminants::CardanoTransactions,
172                SignedEntityTypeDiscriminants::CardanoStakeDistribution,
173            ],
174            aggregate_signature_types: vec![AggregateSignatureType::Concatenation],
175        };
176        let required_capabilities = command.build_required_capabilities();
177
178        assert_eq!(
179            required_capabilities,
180            RequiredAggregatorCapabilities::And(vec![
181                RequiredAggregatorCapabilities::And(vec![
182                    RequiredAggregatorCapabilities::SignedEntityType(
183                        SignedEntityTypeDiscriminants::CardanoTransactions
184                    ),
185                    RequiredAggregatorCapabilities::SignedEntityType(
186                        SignedEntityTypeDiscriminants::CardanoStakeDistribution
187                    ),
188                ]),
189                RequiredAggregatorCapabilities::And(vec![
190                    RequiredAggregatorCapabilities::AggregateSignatureType(
191                        AggregateSignatureType::Concatenation
192                    ),
193                ]),
194            ])
195        );
196    }
197}