warp generated capabilities to provide some comfort functions

Lachstec · Lachstec · Feb 15, 2024 · Feb 15, 2024 · Feb 15, 2024 · Feb 15, 2024
commit fdda741ac89ec8ee87d699fa9f7e5f942622f571
diff --git a/src/client/capabilities.rs b/src/client/capabilities.rs
@@ -0,0 +1,49 @@
+use crate::gen::gnmi::CapabilityResponse;
+use crate::gen::gnmi::ModelData;
+
+
+pub use crate::gen::gnmi::Encoding;
+
+/// Capabilities of a given gNMI Target device.
+///
+/// Contains information about the capabilities that supported by a gNMI Target device.
+#[derive(Debug, Clone)]
+pub struct Capabilities(pub(crate) CapabilityResponse);
+
+impl<'a> Capabilities {
+    /// Retrieve the gNMI Version that the target device supports.
+    pub fn gnmi_version(&'a self) -> &'a str {
+        self.0.g_nmi_version.as_str()
+    }
+
+    /// Check if target device supports a given model.
+    ///
+    /// # Arguments
+    /// - name: Name of the model
+    /// - organization: Organization publishing the model
+    /// - version: Version of the model
+    pub fn supports_model(&self, name: &str, organization: &str, version: &str) -> bool {
+        self.0.supported_models.contains(&ModelData {
+            name: name.to_string(),
+            organization: organization.to_string(),
+            version: version.to_string()
+        })
+    }
+
+    /// Check if a target device supports a given [`Encoding`].
+    ///
+    /// # Arguments
+    /// - encoding: The [`Encoding`] to check for.
+    pub fn supports_encoding(&self, encoding: Encoding) -> bool {
+        let enc: i32 = match encoding {
+            Encoding::Json => 0,
+            Encoding::Bytes => 1,
+            Encoding::Proto => 2,
+            Encoding::Ascii => 3,
+            Encoding::JsonIetf => 4
+        };
+
+        self.0.supported_encodings.contains(&enc)
+    }
+}
+
diff --git a/src/client/client.rs b/src/client/client.rs
@@ -1,7 +1,8 @@
 use crate::auth::AuthService;
 use crate::error::GinmiError;
 use crate::gen::gnmi::g_nmi_client::GNmiClient;
-use crate::gen::gnmi::{CapabilityRequest, CapabilityResponse};
+use crate::gen::gnmi::CapabilityRequest;
+use super::capabilities::Capabilities;
 use http::HeaderValue;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -36,10 +37,10 @@ impl<'a> Client {
     /// let capabilities = client.capabilities().await;
     /// # }
     /// ```
-    pub async fn capabilities(&mut self) -> CapabilityResponse {
+    pub async fn capabilities(&mut self) -> Capabilities {
         let req = CapabilityRequest::default();
         match self.inner.capabilities(req).await {
-            Ok(val) => val.into_inner(),
+            Ok(val) => Capabilities(val.into_inner()),
             Err(e) => panic!("Error getting capabilities: {:?}", e),
         }
     }

diff --git a/src/client/mod.rs b/src/client/mod.rs
@@ -1,3 +1,9 @@
 mod client;
+mod capabilities;
 
-pub use client::{Client, ClientBuilder};
+pub use client::{Client, ClientBuilder};
+
+pub use capabilities::{
+    Capabilities,
+    Encoding
+};
diff --git a/src/lib.rs b/src/lib.rs
@@ -6,8 +6,12 @@ mod auth;
 mod client;
 mod error;
 
-pub use client::{Client, ClientBuilder};
-
+pub use client::{
+    Client,
+    ClientBuilder,
+    Encoding,
+    Capabilities
+};
 pub use error::GinmiError;
 
 pub(crate) mod gen {