Skip to content

Commit

Permalink
fix: simplified mock
Browse files Browse the repository at this point in the history 
Signed-off-by: WoodenMaiden <[email protected]>
  • Loading branch information
@WoodenMaiden
WoodenMaiden committed Oct 28, 2023
1 parent f1fd44c commit f1fe5b6
Showing 1 changed file with 43 additions and 97 deletions.
140 changes: 43 additions & 97 deletions initramfs/src/image.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ impl Image {
&self.tag
}

pub fn export_to_initramfs<H: FileHandler>(
pub fn export_to_initramfs<Handler: FileHandler>(
&self,
init_path: &str,
agent_path: &str,
Expand All @@ -90,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(
<H>::create(&file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?,
<Handler>::create(&file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?,
)
.map_err(|e| anyhow!(e).context("Failed to create gzip encoder"))?;

Expand Down Expand Up @@ -147,10 +147,10 @@ impl Image {
}

let mut init_file =
<H>::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
<Handler>::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
let mut agent_file =
<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)
<Handler>::open(agent_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?;
let mut agent_config_file = <Handler>::open(agent_config_path)
.map_err(|e| anyhow!(e).context("Failed to open agent config file"))?;

let mut init_content = Vec::new();
Expand Down Expand Up @@ -199,69 +199,24 @@ impl Image {
#[cfg(test)]
mod test {
use super::{FileHandler, Image};
use anyhow::{anyhow, Ok};
use anyhow::Ok;
use std::env;
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);

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 flush(&mut self) -> std::io::Result<()> {
std::result::Result::Ok(())
}
}

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(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
std::result::Result::Ok(0)
}

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())
}
}
Expand All @@ -271,15 +226,7 @@ mod test {
where
Self: std::marker::Sized,
{
unsafe {
FILES.push(Self {
contents: vec![],
path: path.to_string(),
});
}

Ok(Self {
contents: vec![],
path: path.to_string(),
})
}
Expand All @@ -288,22 +235,13 @@ mod test {
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());
}
}
}
MockFileHandler::create(path)
}
}

#[test]
pub fn valid_image_name() {

let image1 = Image::new(VALID_IMAGE_NAME);
assert!(image1.is_ok());

Expand All @@ -319,11 +257,36 @@ mod test {
assert_eq!(image2.tag(), "latest");
}

#[test]
pub fn invalid_image_name() {
let image = Image::new(":tag_but_with_no_image");
assert!(image.is_err());
}

#[test]
pub fn test_initramfs_export() {
// setup

// holds the CPIO file
static mut FILE: Vec<u8> = vec![];

impl Write for MockFileHandler {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
unsafe {
if self.path.contains(".img") && FILE.len() < 2 {
FILE.extend_from_slice(buf)
}
};

std::result::Result::Ok(buf.len())
}

fn flush(&mut self) -> std::io::Result<()> {
std::result::Result::Ok(())
}
}

let image = Image::new(VALID_IMAGE_NAME);
assert!(image.is_ok());

let image_filename = format!("{}/init", env::temp_dir().display());
let agent_filename = format!("{}/agent", env::temp_dir().display());
Expand All @@ -333,31 +296,14 @@ mod test {
MockFileHandler::create(&agent_filename).unwrap();
MockFileHandler::create(&agent_config_filename).unwrap();

let image = image.unwrap();
let generated_filename = format!(
"initramfs-{}-{}.img",
image.name.replace('/', "-"),
image.tag
);

// checks
assert!(image
.export_to_initramfs::<MockFileHandler>(
image_filename.as_str(),
agent_filename.as_str(),
agent_config_filename.as_str()
)
.is_ok());

let generated_file = MockFileHandler::open(&generated_filename);

let magic_number = generated_file
.unwrap()
.bytes()
.take(2)
.map(|b| b.unwrap())
.collect::<Vec<u8>>();

assert_eq!(magic_number, &[0x1F, 0x8b]);
assert!(image.unwrap().export_to_initramfs::<MockFileHandler>(
image_filename.as_str(),
agent_filename.as_str(),
agent_config_filename.as_str()
).is_ok());

assert_eq!(unsafe { FILE.clone() }, vec![0x1F, 0x8b]);
}
}

0 comments on commit f1fe5b6

Please sign in to comment.