From b8829bbf221dc5611973d524305ca59aaf10b3d6 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Sun, 10 Mar 2019 17:10:41 -0700
Subject: [PATCH] add identity middleware tests

---
 src/middleware/identity.rs | 66 ++++++++++++++++++++++++++++++++++++++
 src/test.rs                |  7 ++++
 2 files changed, 73 insertions(+)

diff --git a/src/middleware/identity.rs b/src/middleware/identity.rs
index f3ccca932..d0a4146ae 100644
--- a/src/middleware/identity.rs
+++ b/src/middleware/identity.rs
@@ -461,3 +461,69 @@ impl IdentityPolicy for CookieIdentityPolicy {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::http::StatusCode;
+    use crate::test::{self, TestRequest};
+    use crate::{web, App, HttpResponse};
+
+    #[test]
+    fn test_identity() {
+        let mut srv = test::init_service(
+            App::new()
+                .middleware(IdentityService::new(
+                    CookieIdentityPolicy::new(&[0; 32])
+                        .domain("www.rust-lang.org")
+                        .name("actix_auth")
+                        .path("/")
+                        .secure(true),
+                ))
+                .service(web::resource("/index").to(|id: Identity| {
+                    if id.identity().is_some() {
+                        HttpResponse::Created()
+                    } else {
+                        HttpResponse::Ok()
+                    }
+                }))
+                .service(web::resource("/login").to(|id: Identity| {
+                    id.remember("test".to_string());
+                    HttpResponse::Ok()
+                }))
+                .service(web::resource("/logout").to(|id: Identity| {
+                    if id.identity().is_some() {
+                        id.forget();
+                        HttpResponse::Ok()
+                    } else {
+                        HttpResponse::BadRequest()
+                    }
+                })),
+        );
+        let resp =
+            test::call_success(&mut srv, TestRequest::with_uri("/index").to_request());
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let resp =
+            test::call_success(&mut srv, TestRequest::with_uri("/login").to_request());
+        assert_eq!(resp.status(), StatusCode::OK);
+        let c = resp.cookies().next().unwrap().to_owned();
+
+        let resp = test::call_success(
+            &mut srv,
+            TestRequest::with_uri("/index")
+                .cookie(c.clone())
+                .to_request(),
+        );
+        assert_eq!(resp.status(), StatusCode::CREATED);
+
+        let resp = test::call_success(
+            &mut srv,
+            TestRequest::with_uri("/logout")
+                .cookie(c.clone())
+                .to_request(),
+        );
+        assert_eq!(resp.status(), StatusCode::OK);
+        assert!(resp.headers().contains_key(header::SET_COOKIE))
+    }
+}
diff --git a/src/test.rs b/src/test.rs
index 445924006..03700b679 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -11,6 +11,7 @@ use actix_rt::Runtime;
 use actix_server_config::ServerConfig;
 use actix_service::{IntoNewService, NewService, Service};
 use bytes::Bytes;
+use cookie::Cookie;
 use futures::Future;
 
 use crate::config::{AppConfig, AppConfigInner};
@@ -241,6 +242,12 @@ impl TestRequest {
         self
     }
 
+    /// Set cookie for this request
+    pub fn cookie<'a>(mut self, cookie: Cookie<'a>) -> Self {
+        self.req.cookie(cookie);
+        self
+    }
+
     /// Set request payload
     pub fn set_payload<B: Into<Bytes>>(mut self, data: B) -> Self {
         self.req.set_payload(data);