add support to add cookies to the request

master
kirbylife 2023-08-06 00:36:50 -06:00
parent 074ba2bc10
commit edc7582de0
2 changed files with 65 additions and 1 deletions

View File

@ -1,5 +1,7 @@
use http::Method; use http::header::{HeaderName, HeaderValue};
use http::{HeaderMap, Method};
use std::collections::HashMap; use std::collections::HashMap;
use std::str::FromStr;
use url::ParseError; use url::ParseError;
use crate::into_url::IntoUrl; use crate::into_url::IntoUrl;
@ -31,6 +33,7 @@ pub struct Request {
method: Method, method: Method,
params: Vec<(String, String)>, params: Vec<(String, String)>,
headers: HashMap<String, String>, headers: HashMap<String, String>,
cookies: HashMap<String, String>,
} }
impl RequestBase for Request { impl RequestBase for Request {
@ -41,6 +44,7 @@ impl RequestBase for Request {
method: Method::GET, method: Method::GET,
params: Vec::new(), params: Vec::new(),
headers: HashMap::new(), headers: HashMap::new(),
cookies: HashMap::new(),
}) })
} }
@ -55,6 +59,24 @@ impl RequestBase for Request {
self.url.query_pairs_mut().append_pair(key, value); self.url.query_pairs_mut().append_pair(key, value);
} }
let cookies = self
.cookies
.iter()
.map(|item| format!("{}={}", item.0, item.1))
.collect::<Vec<_>>()
.join(";");
if !self.cookies.is_empty() {
self.headers.insert("Cookie".to_string(), cookies);
}
let mut headers = HeaderMap::new();
for (name, header) in self.headers {
let header_name = HeaderName::from_str(name.as_str()).unwrap();
let header_value = HeaderValue::from_str(header.as_str()).unwrap();
headers.insert(header_name, header_value);
}
let resp = (match self.method { let resp = (match self.method {
Method::GET => client.get(self.url.as_str()), Method::GET => client.get(self.url.as_str()),
Method::POST => client.post(self.url.as_str()), Method::POST => client.post(self.url.as_str()),
@ -63,6 +85,7 @@ impl RequestBase for Request {
Method::DELETE => client.delete(self.url.as_str()), Method::DELETE => client.delete(self.url.as_str()),
_ => unimplemented!(), _ => unimplemented!(),
}) })
.headers(headers)
.send() .send()
.unwrap(); .unwrap();
@ -113,4 +136,17 @@ impl Request {
self.params self.params
.push((key.as_ref().to_string(), value.to_string())); .push((key.as_ref().to_string(), value.to_string()));
} }
pub fn add_cookies<S: AsRef<str>, T: ToString>(mut self, cookies: Vec<(S, T)>) -> Self {
self.cookies = cookies
.iter()
.map(|item| (item.0.as_ref().to_string(), item.1.to_string()))
.collect();
self
}
pub fn add_cookie<S: AsRef<str>, T: ToString>(&mut self, key: S, value: T) {
self.cookies
.insert(key.as_ref().to_string(), value.to_string());
}
} }

View File

@ -1,6 +1,7 @@
use http::Method; use http::Method;
use http::StatusCode; use http::StatusCode;
use raspa::request::{ParameterType, Request, RequestBase}; use raspa::request::{ParameterType, Request, RequestBase};
use raspa::response::ResponseBase;
use raspa::selector::{Selector, SelectorBase}; use raspa::selector::{Selector, SelectorBase};
use serde_json::Value; use serde_json::Value;
// use std::collections::HashMap; // use std::collections::HashMap;
@ -172,3 +173,30 @@ fn custom_user_agent() {
Some("test user agent") Some("test user agent")
); );
} }
#[test]
fn simple_cookie_test() {
let mut req = Request::new("https://httpbin.org/cookies").unwrap();
req.add_cookie("test_cookie", "100");
let resp: Value = req.launch().to_json().expect("Weird JSON parsed fail");
assert_eq!(resp["cookies"]["test_cookie"].as_str(), Some("100"));
}
#[test]
fn multiple_cookies_test() {
let resp: Value = Request::new("https://httpbin.org/cookies")
.unwrap()
.add_cookies(vec![
("cookie_1", 100),
("cookie_2", 200),
("cookie_3", 300),
])
.launch()
.to_json()
.expect("Weird JSON parsed fail");
assert_eq!(resp["cookies"]["cookie_1"].as_str(), Some("100"));
assert_eq!(resp["cookies"]["cookie_2"].as_str(), Some("200"));
assert_eq!(resp["cookies"]["cookie_3"].as_str(), Some("300"));
}