EduardoRodriguesF · EduardoRodriguesF · May 28, 2024 · May 19, 2024 · May 19, 2024 · May 19, 2024
diff --git a/src/action/env.rs b/src/action/env.rs
@@ -1,6 +1,10 @@
+use core::panic;
 use std::process::ExitCode;
 
-use crate::{cli::EnvCmd as Cmd, Ctx, Env, QuartzResult, StateField};
+use crate::{
+    cli::{EnvCmd as Cmd, HeaderCmd},
+    Ctx, Env, PairMap, QuartzResult, StateField,
+};
 use colored::Colorize;
 
 #[derive(clap::Args, Debug)]
@@ -31,6 +35,12 @@ pub fn cmd(ctx: &mut Ctx, command: Cmd) -> QuartzResult {
         Cmd::Use(args) => switch(ctx, args)?,
         Cmd::Ls => ls(ctx),
         Cmd::Rm(args) => rm(ctx, args),
+        Cmd::Header { command } => match command {
+            HeaderCmd::Set { header } => header_set(ctx, header)?,
+            HeaderCmd::Ls => header_ls(ctx)?,
+            HeaderCmd::Rm { key } => header_rm(ctx, key)?,
+            HeaderCmd::Get { key } => header_get(ctx, key)?,
+        },
     };
 
     Ok(())
@@ -136,3 +146,33 @@ pub fn print(ctx: &Ctx) {
             .unwrap_or("default".into())
     );
 }
+pub fn header_set(ctx: &Ctx, args: Vec<String>) -> QuartzResult {
+    let mut env = ctx.require_env();
+    for header in args {
+        env.headers.set(&header);
+    }
+    env.update(ctx)?;
+    Ok(())
+}
+pub fn header_ls(ctx: &Ctx) -> QuartzResult {
+    let env = ctx.require_env();
+    print!("{}", env.headers);
+    Ok(())
+}
+pub fn header_rm(ctx: &Ctx, keys: Vec<String>) -> QuartzResult {
+    for key in keys {
+        let mut env = ctx.require_env();
+        env.headers.remove(&key);
+        env.update(ctx)?;
+    }
+    Ok(())
+}
+pub fn header_get(ctx: &Ctx, key: String) -> QuartzResult {
+    let env = ctx.require_env();
+    let value = env
+        .headers
+        .get(&key)
+        .unwrap_or_else(|| panic!("no header named {key} found"));
+    println!("{value}");
+    Ok(())
+}
diff --git a/src/action/send.rs b/src/action/send.rs
@@ -7,6 +7,7 @@ use crate::{
 use chrono::Utc;
 use hyper::{
     body::{Bytes, HttpBody},
+    header::{HeaderName, HeaderValue},
     Body, Client, Uri,
 };
 use std::path::{Path, PathBuf};
@@ -93,11 +94,17 @@ pub async fn cmd(ctx: &Ctx, mut args: Args) -> QuartzResult {
     let mut res: hyper::Response<Body>;
 
     loop {
-        let req = endpoint
+        let mut req = endpoint
             // TODO: Find a way around this clone
             .clone()
             .into_request()
             .unwrap_or_else(|_| panic!("malformed request"));
+        for (key, val) in env.headers.iter() {
+            if !endpoint.headers.contains_key(key) {
+                req.headers_mut()
+                    .insert(HeaderName::from_str(key)?, HeaderValue::from_str(val)?);
+            }
+        }
 
         entry.message(&req);
         if let Some(ref body) = body {

diff --git a/src/cli.rs b/src/cli.rs
@@ -231,6 +231,10 @@ pub enum EnvCmd {
     /// Delete a environment
     #[command(name = "rm", alias = "remove")]
     Rm(action::env::RmArgs),
+    Header {
+        #[command(subcommand)]
+        command: HeaderCmd,
+    },
 }
 
 #[derive(Debug, Subcommand)]

diff --git a/src/endpoint.rs b/src/endpoint.rs
@@ -51,6 +51,16 @@ impl PairMap<'_> for Query {
 #[derive(Default, Debug, Serialize, Deserialize, Clone)]
 pub struct Headers(pub HashMap<String, String>);
 
+impl Headers {
+    pub fn parse(file_content: &str) -> Self {
+        let mut headers = Headers::default();
+        for header in file_content.lines().filter(|line| !line.is_empty()) {
+            headers.set(header);
+        }
+        headers
+    }
+}
+
 impl Deref for Headers {
     type Target = HashMap<String, String>;
 

diff --git a/src/env.rs b/src/env.rs
@@ -8,7 +8,7 @@ use std::{
 
 use serde::{Deserialize, Serialize};
 
-use crate::{cookie::CookieJar, Ctx, PairMap};
+use crate::{cookie::CookieJar, endpoint::Headers, Ctx, PairMap};
 
 #[derive(Default, Debug, Clone, Serialize, Deserialize)]
 pub struct Variables(pub HashMap<String, String>);
@@ -61,13 +61,15 @@ impl Variables {
 pub struct Env {
     pub name: String,
     pub variables: Variables,
+    pub headers: Headers,
 }
 
 impl Default for Env {
     fn default() -> Self {
         Self {
             name: String::from("default"),
             variables: Variables::default(),
+            headers: Headers::default(),
         }
     }
 }
@@ -100,10 +102,18 @@ impl Env {
             .write(true)
             .truncate(true)
             .open(self.dir(ctx).join("variables"))?;
+        let mut headers_file = std::fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(true)
+            .open(self.dir(ctx).join("headers"))?;
 
         if !self.variables.is_empty() {
             var_file.write_all(format!("{}", self.variables).as_bytes())?;
         }
+        if !self.headers.0.is_empty() {
+            headers_file.write_all(format!("{}", self.headers).as_bytes())?;
+        }
 
         Ok(())
     }
@@ -119,6 +129,9 @@ impl Env {
         if let Ok(var_contents) = std::fs::read_to_string(env.dir(ctx).join("variables")) {
             env.variables = Variables::parse(&var_contents);
         }
+        if let Ok(header_contents) = std::fs::read_to_string(env.dir(ctx).join("headers")) {
+            env.headers = Headers::parse(&header_contents);
+        }
 
         Ok(env)
     }

diff --git a/tests/it/env.rs b/tests/it/env.rs
@@ -116,3 +116,31 @@ fn it_cannot_create_duplicate() -> TestResult {
 
     Ok(())
 }
+
+#[test]
+fn it_create_headers() -> TestResult {
+    let quartz = Quartz::preset_empty_project()?;
+    quartz.cmd(&["env", "header", "set", "Header1: Value1"])?;
+    quartz.cmd(&["env", "header", "set", "Header2: Value2"])?;
+
+    let output = quartz.cmd(&["env", "header", "get", "Header1"])?;
+    assert_eq!(output.stdout.trim(), "Value1");
+
+    let output = quartz.cmd(&["env", "header", "get", "Header2"])?;
+    assert_eq!(output.stdout.trim(), "Value2");
+    Ok(())
+}
+#[test]
+fn it_delete_header() -> TestResult {
+    let quartz = Quartz::preset_empty_project()?;
+    quartz.cmd(&["env", "header", "set", "Header1: Value1"])?;
+
+    let output = quartz.cmd(&["env", "header", "get", "Header1"])?;
+    assert_eq!(output.stdout.trim(), "Value1");
+
+    quartz.cmd(&["env", "header", "rm", "Header1"])?;
+    let output = quartz.cmd(&["env", "header", "get", "Header1"])?;
+    assert!(!output.status.success());
+
+    Ok(())
+}