From 1cf010b5a0ffe8b2b573375e1f092fcd891b4d77 Mon Sep 17 00:00:00 2001
From: kirbylife <gabriel13m@gmail.com>
Date: Sat, 30 Jan 2021 21:11:54 -0600
Subject: [PATCH] add 'attr' method in the SelectorBase

---
 src/lib.rs          | 36 -------------------------
 src/selector.rs     | 10 ++++++-
 tests/unit_tests.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 37 deletions(-)
 create mode 100644 tests/unit_tests.rs

diff --git a/src/lib.rs b/src/lib.rs
index 51b9122..0e5b335 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,39 +1,3 @@
 pub mod request;
 pub mod response;
 pub mod selector;
-
-#[cfg(test)]
-mod tests {
-    use crate::request::{Request, RequestBase};
-    use crate::selector::{Selector, SelectorBase};
-    use http::StatusCode;
-
-    #[test]
-    fn plain_text_selector() {
-        let html = "
-<html>
-  <body>
-    <h1>hello world</h1>
-    <p id='text'>good bye</p>
-    <a>simple text</a>
-  </body>
-</html>
-"
-        .to_string();
-        let sel = Selector::from_html(html);
-        assert_eq!(sel.css::<Selector>("h1")[0].html(), "<h1>hello world</h1>");
-        assert_eq!(sel.css::<Selector>("#text")[0].content(), "good bye");
-        assert_eq!(
-            sel.css_once::<Selector>("body > a").unwrap().content(),
-            "simple text"
-        );
-    }
-
-    #[test]
-    fn simple_request() {
-        let req: Request = RequestBase::new("https://httpbin.org/").unwrap();
-        let resp = req.launch();
-        assert_eq!(resp.status_code, StatusCode::OK);
-        assert!(resp.css::<Selector>("h2")[0].html().contains("httpbin.org"));
-    }
-}
diff --git a/src/selector.rs b/src/selector.rs
index 7b40b02..989718f 100644
--- a/src/selector.rs
+++ b/src/selector.rs
@@ -20,12 +20,20 @@ pub trait SelectorBase {
 
     fn content(&self) -> String {
         let html = nipper::Document::from(self.html().as_str());
-        html.select("*")
+        html.select("body > *")
             .iter()
             .map(|element| element.text().to_string())
             .last()
             .unwrap()
     }
+
+    fn attr(&self, attribute: &'static str) -> Option<String> {
+        let html = nipper::Document::from(self.html().as_str());
+        match html.select("body > *").attr(attribute) {
+            Some(text) => Some(format!("{}", text)),
+            None => None,
+        }
+    }
 }
 
 #[derive(Debug)]
diff --git a/tests/unit_tests.rs b/tests/unit_tests.rs
new file mode 100644
index 0000000..b96df90
--- /dev/null
+++ b/tests/unit_tests.rs
@@ -0,0 +1,64 @@
+use http::StatusCode;
+use raspa::request::{Request, RequestBase};
+use raspa::selector::{Selector, SelectorBase};
+
+#[test]
+fn plain_text_selector() {
+    let html = "
+<html>
+  <body>
+    <h1>hello world</h1>
+    <p id='text'>good bye</p>
+    <a>simple text</a>
+  </body>
+</html>
+"
+    .to_string();
+    let sel = Selector::from_html(html);
+    assert_eq!(sel.css::<Selector>("h1")[0].html(), "<h1>hello world</h1>");
+    assert_eq!(sel.css::<Selector>("#text")[0].content(), "good bye");
+    assert_eq!(
+        sel.css_once::<Selector>("body > a").unwrap().content(),
+        "simple text"
+    );
+}
+
+#[test]
+fn simple_request() {
+    let req: Request = RequestBase::new("https://httpbin.org/").unwrap();
+    let resp = req.launch();
+    assert_eq!(resp.status_code, StatusCode::OK);
+    assert!(resp.css::<Selector>("h2")[0].html().contains("httpbin.org"));
+}
+
+#[test]
+fn complex_selectors() {
+    let html = "
+<html>
+  <body>
+    <p id='text'>good bye</p>
+    <a href='http://google.com'>simple text</a>
+    <ul>
+        <li class='item' id='item-1'>1</li>
+        <li class='item' id='item-2'>2</li>
+        <li class='item' id='item-3'>3</li>
+    </ul>
+  </body>
+</html>
+"
+    .to_string();
+    let sel = Selector::from_html(html);
+    assert_eq!(
+        sel.css_once::<Selector>("p").unwrap().attr("id").unwrap(),
+        "text"
+    );
+    assert_eq!(
+        sel.css::<Selector>("a")[0].attr("href").unwrap(),
+        "http://google.com"
+    );
+    for node in sel.css::<Selector>("ul li").iter() {
+        let text = node.content();
+        assert_eq!(node.attr("class").unwrap(), "item");
+        assert!(node.attr("id").unwrap().contains(&text));
+    }
+}