test: split composition tests

gusinacio · gusinacio · commit 77aaa773aa01 · 2024-11-22T13:50:33.000-06:00
Signed-off-by: Gustavo Inacio &lt;gustavo@semiotic.ai&gt;
diff --git a/crates/service/src/middleware/auth.rs b/crates/service/src/middleware/auth.rs
@@ -19,6 +19,7 @@ mod tests {
     use reqwest::{header, StatusCode};
     use sqlx::PgPool;
     use tap_core::{manager::Manager, receipt::checks::CheckList};
+    use tokio::time::sleep;
     use tower::{Service, ServiceBuilder, ServiceExt};
     use tower_http::auth::AsyncRequireAuthorizationLayer;
 
@@ -27,80 +28,111 @@ mod tests {
     use test_assets::{create_signed_receipt, TAP_EIP712_DOMAIN};
 
     const ALLOCATION_ID: Address = address!("deadbeefcafebabedeadbeefcafebabedeadbeef");
+    const BEARER_TOKEN: &str = "test";
 
-    async fn handle(_: Request<Body>) -> anyhow::Result<Response<Body>> {
-        Ok(Response::new(Body::default()))
-    }
-
-    #[sqlx::test(migrations = "../../migrations")]
-    async fn test_middleware_composition(pgpool: PgPool) {
-        let token = "test".to_string();
+    async fn service(
+        pgpool: PgPool,
+    ) -> impl Service<Request<Body>, Response = Response<Body>, Error = impl std::fmt::Debug> {
         let context = IndexerTapContext::new(pgpool.clone(), TAP_EIP712_DOMAIN.clone()).await;
         let tap_manager = Box::leak(Box::new(Manager::new(
             TAP_EIP712_DOMAIN.clone(),
             context,
             CheckList::empty(),
         )));
+
+        let registry = prometheus::Registry::new();
         let metric = Box::leak(Box::new(
-            prometheus::register_counter_vec!(
+            prometheus::register_counter_vec_with_registry!(
                 "merge_checks_test",
                 "Failed queries to handler",
-                &["deployment"]
+                &["deployment"],
+                registry,
             )
             .unwrap(),
         ));
-        let free_query = Bearer::new(&token);
+        let free_query = Bearer::new(BEARER_TOKEN);
         let tap_auth = auth::tap_receipt_authorize(tap_manager, metric);
         let authorize_requests = free_query.or(tap_auth);
 
         let authorization_middleware = AsyncRequireAuthorizationLayer::new(authorize_requests);
 
         let mut service = ServiceBuilder::new()
             .layer(authorization_middleware)
-            .service_fn(handle);
+            .service_fn(|_: Request<Body>| async {
+                Ok::<_, anyhow::Error>(Response::new(Body::default()))
+            });
 
-        let handle = service.ready().await.unwrap();
+        service.ready().await.unwrap();
+        service
+    }
 
+    #[sqlx::test(migrations = "../../migrations")]
+    async fn test_composition_header_valid(pgpool: PgPool) {
+        let mut service = service(pgpool.clone()).await;
         // should allow queries that contains the free token
         // if the token does not match, return payment required
         let mut req = Request::new(Default::default());
         req.headers_mut().insert(
             header::AUTHORIZATION,
-            format!("Bearer {token}").parse().unwrap(),
+            format!("Bearer {}", BEARER_TOKEN).parse().unwrap(),
         );
-        let res = handle.call(req).await.unwrap();
+        let res = service.call(req).await.unwrap();
         assert_eq!(res.status(), StatusCode::OK);
+    }
+
+    #[sqlx::test(migrations = "../../migrations")]
+    async fn test_composition_header_invalid(pgpool: PgPool) {
+        let mut service = service(pgpool.clone()).await;
 
         // if the token exists but is wrong, try the receipt
         let mut req = Request::new(Default::default());
         req.headers_mut()
             .insert(header::AUTHORIZATION, "Bearer wrongtoken".parse().unwrap());
-        let res = handle.call(req).await.unwrap();
+        let res = service.call(req).await.unwrap();
         // we return the error from tap
         assert_eq!(res.status(), StatusCode::PAYMENT_REQUIRED);
+    }
+
+    #[sqlx::test(migrations = "../../migrations")]
+    async fn test_composition_with_receipt(pgpool: PgPool) {
+        let mut service = service(pgpool.clone()).await;
 
         let receipt = create_signed_receipt(ALLOCATION_ID, 1, 1, 1).await;
 
         // check with receipt
         let mut req = Request::new(Default::default());
         req.extensions_mut().insert(receipt);
-        let res = handle.call(req).await.unwrap();
+        let res = service.call(req).await.unwrap();
         assert_eq!(res.status(), StatusCode::OK);
 
-        // todo make this sleep better
-        tokio::time::sleep(Duration::from_millis(100)).await;
-
         // verify receipts
-        let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
-            .fetch_all(&pgpool)
-            .await
-            .unwrap();
-        assert_eq!(result.len(), 1);
-
-        // if it has neither, should return unauthorized
-        // check no headers
+        if tokio::time::timeout(Duration::from_secs(1), async {
+            loop {
+                let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
+                    .fetch_all(&pgpool)
+                    .await
+                    .unwrap();
+
+                if result.is_empty() {
+                    sleep(Duration::from_millis(50)).await;
+                } else {
+                    break;
+                }
+            }
+        })
+        .await
+        .is_err()
+        {
+            panic!("Timeout assertion");
+        }
+    }
+
+    #[sqlx::test(migrations = "../../migrations")]
+    async fn test_composition_without_header_or_receipt(pgpool: PgPool) {
+        let mut service = service(pgpool.clone()).await;
+        // if it has neither, should return payment required
         let req = Request::new(Default::default());
-        let res = handle.call(req).await.unwrap();
+        let res = service.call(req).await.unwrap();
         assert_eq!(res.status(), StatusCode::PAYMENT_REQUIRED);
     }
 }
