fix(test): create trait to test file compression

Signed-off-by: WoodenMaiden <[email protected]>

initramfs/Cargo.toml

@@ -21,6 +21,3 @@ bytes = "1.4.0"
 libflate = "0.1"
 cpio = "0.2.2"
 openssl = { version = "0.10", features = ["vendored"] }
-
-[dev-dependencies]
-infer = "0.15.0"

initramfs/src/image.rs

@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     fs::File,
-    io::{Cursor, Read},
+    io::{Cursor, Read, Write},
 };
 
 use anyhow::{anyhow, Result};
@@ -12,6 +12,26 @@ use log::debug;
 use serde::Deserialize;
 use tar::Archive;
 
+/// Trait to abstract File reading
+pub trait FileHandler: Write + Read {
+    fn create(path: &str) -> Result<Self>
+    where
+        Self: std::marker::Sized;
+    fn open(path: &str) -> Result<Self>
+    where
+        Self: std::marker::Sized;
+}
+
+impl FileHandler for File {
+    fn create(path: &str) -> Result<Self> {
+        File::create(path).map_err(|e| anyhow!(e).context("Failed to create file"))
+    }
+
+    fn open(path: &str) -> Result<Self> {
+        File::open(path).map_err(|e| anyhow!(e).context("Failed to open file"))
+    }
+}
+
 #[derive(Deserialize, Debug, Clone)]
 #[allow(non_snake_case)]
 pub struct LayerMetadata {
@@ -61,7 +81,7 @@ impl Image {
         &self.tag
     }
 
-    pub fn export_to_initramfs(
+    pub fn export_to_initramfs<H: FileHandler>(
         &self,
         init_path: &str,
         agent_path: &str,
@@ -70,7 +90,7 @@ impl Image {
         // Write the cpio to disk
         let file_name = format!("initramfs-{}-{}.img", self.name.replace('/', "-"), self.tag);
         let archive = Encoder::new(
-            File::create(file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?,
+            <H>::create(&file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?,
         )
         .map_err(|e| anyhow!(e).context("Failed to create gzip encoder"))?;
 
@@ -127,10 +147,10 @@ impl Image {
         }
 
         let mut init_file =
-            File::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
+            <H>::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
         let mut agent_file =
-            File::open(agent_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
-        let mut agent_config_file = File::open(agent_config_path)
+            <H>::open(agent_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
+        let mut agent_config_file = <H>::open(agent_config_path)
             .map_err(|e| anyhow!(e).context("Failed to open agent config file"))?;
 
         let mut init_content = Vec::new();
@@ -178,23 +198,109 @@ impl Image {
 
 #[cfg(test)]
 mod test {
+    use super::{FileHandler, Image};
+    use anyhow::{anyhow, Ok};
     use std::env;
-    use std::io;
-    use std::io::Read;
-    use std::{error::Error, result::Result};
+    use std::io::{Read, Write};
+    use std::vec::Drain;
+
+    const VALID_IMAGE_NAME: &str = "my_awesome_img:14md35";
+    const VALID_IMAGE_NAME_FROM_HUB: &str = "bitnami/mongodb:latest";
+
+    #[derive(Debug, Clone)]
+    struct MockFileHandler {
+        contents: Vec<u8>,
+        path: String,
+    }
+
+    // holds all the files created by the tests
+    static mut FILES: Vec<MockFileHandler> = vec![];
+
+    impl Write for MockFileHandler {
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            self.contents.extend_from_slice(buf);
 
-    use super::Image;
-    use std::fs::{File, OpenOptions};
+            unsafe {
+                FILES.iter_mut().for_each(|f| {
+                    if f.path == self.path {
+                        f.contents.extend_from_slice(buf);
+                    }
+                });
+            }
+
+            std::result::Result::Ok(buf.len())
+        }
 
-    fn touch(path: &str) -> Result<&str, Box<dyn Error>> {
-        match OpenOptions::new().create(true).write(true).open(path) {
-            Ok(_) => Ok(path),
-            Err(e) => Err(Box::new(e)),
+        fn flush(&mut self) -> std::io::Result<()> {
+            std::result::Result::Ok(())
         }
     }
 
-    const VALID_IMAGE_NAME: &str = "my_awesome_img:14md35";
-    const VALID_IMAGE_NAME_FROM_HUB: &str = "bitnami/mongodb:latest";
+    impl Read for MockFileHandler {
+        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+            let length = if buf.len() > self.contents.len() {
+                self.contents.len()
+            } else {
+                buf.len()
+            };
+
+            let drained = &self.contents.drain(..length);
+            let read = <&Drain<'_, u8>>::clone(&drained).as_slice();
+
+            for i in 0..length {
+                buf[i] = if i < length { read[i] } else { 0 }
+            }
+
+            std::result::Result::Ok(length)
+        }
+
+        fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
+            loop {
+                let nb = self.read(buf)?;
+                if nb == 0 {
+                    break;
+                }
+                buf.extend_from_slice(&self.contents);
+            }
+            std::result::Result::Ok(buf.len())
+        }
+    }
+
+    impl FileHandler for MockFileHandler {
+        fn create(path: &str) -> anyhow::Result<Self>
+        where
+            Self: std::marker::Sized,
+        {
+            unsafe {
+                FILES.push(Self {
+                    contents: vec![],
+                    path: path.to_string(),
+                });
+            }
+
+            Ok(Self {
+                contents: vec![],
+                path: path.to_string(),
+            })
+        }
+
+        fn open(path: &str) -> anyhow::Result<Self>
+        where
+            Self: std::marker::Sized,
+        {
+            unsafe { 
+                let mut iter = FILES.iter_mut();
+
+                loop {
+                    let file = iter.next().ok_or(anyhow!("No file found"))?;
+
+                    if file.path == path {
+                        return Ok(file.clone());
+                    }
+                }
+            }
+        }
+    }
 
     #[test]
     pub fn valid_image_name() {
@@ -223,11 +329,9 @@ mod test {
         let agent_filename = format!("{}/agent", env::temp_dir().display());
         let agent_config_filename = format!("{}/agent_config", env::temp_dir().display());
 
-        let files = [
-            touch(&image_filename).unwrap(),
-            touch(&agent_filename).unwrap(),
-            touch(&agent_config_filename).unwrap(),
-        ];
+        MockFileHandler::create(&image_filename).unwrap();
+        MockFileHandler::create(&agent_filename).unwrap();
+        MockFileHandler::create(&agent_config_filename).unwrap();
 
         let image = image.unwrap();
         let generated_filename = format!(
@@ -238,17 +342,14 @@ mod test {
 
         // checks
         assert!(image
-            .export_to_initramfs(
+            .export_to_initramfs::<MockFileHandler>(
                 image_filename.as_str(),
                 agent_filename.as_str(),
                 agent_config_filename.as_str()
             )
             .is_ok());
 
-        let generated_file = File::open(&generated_filename);
-
-        assert!(std::path::Path::new(&generated_filename).exists());
-        assert!(generated_file.is_ok());
+        let generated_file = MockFileHandler::open(&generated_filename);
 
         let magic_number = generated_file
             .unwrap()
@@ -257,12 +358,6 @@ mod test {
             .map(|b| b.unwrap())
             .collect::<Vec<u8>>();
 
-        let kind = infer::get_from_path(&generated_filename).unwrap().unwrap();
-        assert_eq!(kind.matcher_type(), infer::MatcherType::Archive);
         assert_eq!(magic_number, &[0x1F, 0x8b]);
-
-        // clean
-        files.iter().for_each(|f| std::fs::remove_file(f).unwrap());
-        std::fs::remove_file(&generated_filename).unwrap();
     }
 }

initramfs/src/main.rs

@@ -1,3 +1,5 @@
+use std::fs::File;
+
 use anyhow::anyhow;
 use clap::Parser;
 use env_logger::Env;
@@ -58,7 +60,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     info!("Writing  to disk ...");
     image
-        .export_to_initramfs(&args.init, &args.agent, &args.agent_config)
+        .export_to_initramfs::<File>(&args.init, &args.agent, &args.agent_config)
         .map_err(|e| anyhow!(e).context("Failed to write filesystem to disk"))?;
     info!("Writing done!");