Skip to content

Commit

Permalink
Add encode implementations for SocketAddr and IpAddr
Browse files Browse the repository at this point in the history
  • Loading branch information
rnijveld committed Sep 8, 2022
1 parent 16993b6 commit c8d94d6
Showing 1 changed file with 89 additions and 0 deletions.
89 changes: 89 additions & 0 deletions src/encoding/text.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ use crate::registry::{Registry, Unit};
use std::borrow::Cow;
use std::collections::HashMap;
use std::io::Write;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::ops::Deref;

pub use prometheus_client_derive_text_encode::*;
Expand Down Expand Up @@ -221,6 +222,52 @@ impl Encode for () {
}
}

impl Encode for Ipv4Addr {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
writer.write_all(self.to_string().as_bytes())?;
Ok(())
}
}

impl Encode for Ipv6Addr {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
writer.write_all(self.to_string().as_bytes())?;
Ok(())
}
}

impl Encode for IpAddr {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
match self {
IpAddr::V4(v4) => v4.encode(writer),
IpAddr::V6(v6) => v6.encode(writer),
}
}
}

impl Encode for SocketAddrV4 {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
writer.write_all(self.to_string().as_bytes())?;
Ok(())
}
}

impl Encode for SocketAddrV6 {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
writer.write_all(self.to_string().as_bytes())?;
Ok(())
}
}

impl Encode for SocketAddr {
fn encode(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
match self {
SocketAddr::V4(v4) => v4.encode(writer),
SocketAddr::V6(v6) => v6.encode(writer),
}
}
}

/// Helper type for [`EncodeMetric`], see [`EncodeMetric::encode`].
///
// `Encoder` does not take a trait parameter for `writer` and `labels` because
Expand Down Expand Up @@ -720,6 +767,48 @@ mod tests {
parse_with_python_client(String::from_utf8(encoded).unwrap());
}

#[test]
fn encode_socketaddr() {
let mut registry = Registry::default();
let family = Family::<Vec<(&str, SocketAddr)>, Counter>::default();
registry.register("my_addr", "My socket address", family.clone());

let addr: SocketAddr = "127.0.0.1:80".parse().unwrap();
family.get_or_create(&vec![("address", addr)]).inc();

let mut encoded = Vec::new();
encode(&mut encoded, &registry).unwrap();

let expected = "# HELP my_addr My socket address.\n".to_owned()
+ "# TYPE my_addr counter\n"
+ "my_addr_total{address=\"127.0.0.1:80\"} 1\n"
+ "# EOF\n";
assert_eq!(expected, String::from_utf8(encoded.clone()).unwrap());

parse_with_python_client(String::from_utf8(encoded).unwrap());
}

#[test]
fn encode_ipaddr() {
let mut registry = Registry::default();
let family = Family::<Vec<(&str, IpAddr)>, Counter>::default();
registry.register("my_addr", "My IP address", family.clone());

let addr: IpAddr = "::1".parse().unwrap();
family.get_or_create(&vec![("address", addr)]).inc();

let mut encoded = Vec::new();
encode(&mut encoded, &registry).unwrap();

let expected = "# HELP my_addr My IP address.\n".to_owned()
+ "# TYPE my_addr counter\n"
+ "my_addr_total{address=\"::1\"} 1\n"
+ "# EOF\n";
assert_eq!(expected, String::from_utf8(encoded.clone()).unwrap());

parse_with_python_client(String::from_utf8(encoded).unwrap());
}

#[test]
fn encode_counter_family_with_prefix_with_label() {
let mut registry = Registry::default();
Expand Down

0 comments on commit c8d94d6

Please sign in to comment.