From ada57734176081e109e748f8b3b530cca973d097 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Thu, 8 Jul 2021 21:01:17 +0300
Subject: [PATCH] Convenience TextEncoder functions for working with strings

Signed-off-by: Aleksey Kladov <aleksey.kladov@gmail.com>
---
 src/encoder/text.rs | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/src/encoder/text.rs b/src/encoder/text.rs
index 17516bc0..4a70536f 100644
--- a/src/encoder/text.rs
+++ b/src/encoder/text.rs
@@ -2,6 +2,7 @@
 
 use std::borrow::Cow;
 use std::io::Write;
+use std::mem;
 
 use crate::errors::Result;
 use crate::histogram::BUCKET_LABEL;
@@ -25,6 +26,22 @@ impl TextEncoder {
     pub fn new() -> TextEncoder {
         TextEncoder
     }
+    pub fn encode_utf8(
+        &self,
+        metric_families: &[MetricFamily],
+        buf: &mut String,
+    ) -> Result<()> {
+        let mut bytes = mem::take(buf).into_bytes();
+        self.encode(metric_families, &mut bytes)?;
+        let text = String::from_utf8(bytes).unwrap_or_else(|_| unreachable!());
+        *buf = text;
+        Ok(())
+    }
+    pub fn encode_to_string(&self, metric_families: &[MetricFamily]) -> Result<String> {
+        let mut buf = String::new();
+        self.encode_utf8(metric_families, &mut buf)?;
+        Ok(buf)
+    }
 }
 
 impl Encoder for TextEncoder {
@@ -395,4 +412,25 @@ test_summary_count 5
 "##;
         assert_eq!(ans, str::from_utf8(writer.as_slice()).unwrap());
     }
+
+    #[test]
+    fn test_text_encoder_to_string() {
+        let counter_opts = Opts::new("test_counter", "test help")
+            .const_label("a", "1")
+            .const_label("b", "2");
+        let counter = Counter::with_opts(counter_opts).unwrap();
+        counter.inc();
+
+        let mf = counter.collect();
+
+        let encoder = TextEncoder::new();
+        let txt = encoder.encode_to_string(&mf);
+        let txt = txt.unwrap();
+
+        let counter_ans = r##"# HELP test_counter test help
+# TYPE test_counter counter
+test_counter{a="1",b="2"} 1
+"##;
+        assert_eq!(counter_ans, txt.as_str());
+    }
 }