Fix flakey ratelimit test

src/tests/krate/publish.rs

@@ -1,14 +1,19 @@
-use crate::builders::{CrateBuilder, DependencyBuilder, PublishBuilder};
 use crate::new_category;
 use crate::util::insta::assert_yaml_snapshot;
 use crate::util::{RequestHelper, TestApp};
-use crates_io::controllers::krate::publish::{
-    missing_metadata_error_message, MISSING_RIGHTS_ERROR_MESSAGE,
+use crate::{
+    builders::{CrateBuilder, DependencyBuilder, PublishBuilder},
+    util::TestDatabase,
 };
+use chrono::{NaiveDateTime, Utc};
 use crates_io::models::krate::MAX_NAME_LENGTH;
 use crates_io::rate_limiter::LimitedAction;
 use crates_io::schema::{api_tokens, emails, versions_published_by};
 use crates_io::views::GoodCrate;
+use crates_io::{
+    controllers::krate::publish::{missing_metadata_error_message, MISSING_RIGHTS_ERROR_MESSAGE},
+    schema::{publish_limit_buckets, publish_rate_overrides},
+};
 use crates_io_tarball::TarballBuilder;
 use diesel::{delete, update, ExpressionMethods, QueryDsl, RunQueryDsl};
 use flate2::Compression;
@@ -1006,38 +1011,164 @@ fn tarball_bigger_than_max_upload_size() {
 }
 
 #[test]
-fn publish_new_crate_rate_limited() {
+fn publish_new_crate_ratelimit_hit() {
     let (app, anon, _, token) = TestApp::full()
         .with_rate_limit(LimitedAction::PublishNew, Duration::from_millis(500), 1)
+        .with_database(TestDatabase::TestPool)
         .with_token();
 
-    // Upload a new crate
+    // Set up the database so it'll think we've massively ratelimited ourselves
+    app.db(|conn| {
+        // Ratelimit bucket should next refill in about a year
+        let far_future = Utc::now().naive_utc() + Duration::from_secs(60 * 60 * 24 * 365);
+        diesel::insert_into(publish_limit_buckets::table)
+            .values((
+                publish_limit_buckets::user_id.eq(token.as_model().user_id),
+                publish_limit_buckets::action.eq(LimitedAction::PublishNew),
+                publish_limit_buckets::tokens.eq(0),
+                publish_limit_buckets::last_refill.eq(far_future),
+            ))
+            .execute(conn)
+            .expect("Failed to set fake ratelimit")
+    });
+
+    let crate_to_publish = PublishBuilder::new("rate_limited", "1.0.0");
+    token
+        .publish_crate(crate_to_publish)
+        .assert_rate_limited(LimitedAction::PublishNew);
+
+    assert_eq!(app.stored_files().len(), 0);
+
+    let response = anon.get::<()>("/api/v1/crates/rate_limited");
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
+
+#[test]
+fn publish_new_crate_ratelimit_expires() {
+    let (app, anon, _, token) = TestApp::full()
+        .with_rate_limit(LimitedAction::PublishNew, Duration::from_millis(500), 1)
+        .with_database(TestDatabase::TestPool)
+        .with_token();
+
+    // Set up the database so it'll think we've massively ratelimited ourselves
+    app.db(|conn| {
+        // Ratelimit bucket should next refill right now!
+        let just_now = Utc::now().naive_utc() - Duration::from_millis(500);
+        diesel::insert_into(publish_limit_buckets::table)
+            .values((
+                publish_limit_buckets::user_id.eq(token.as_model().user_id),
+                publish_limit_buckets::action.eq(LimitedAction::PublishNew),
+                publish_limit_buckets::tokens.eq(0),
+                publish_limit_buckets::last_refill.eq(just_now),
+            ))
+            .execute(conn)
+            .expect("Failed to set fake ratelimit")
+    });
+
+    let crate_to_publish = PublishBuilder::new("rate_limited", "1.0.0");
+    token.publish_crate(crate_to_publish).good();
+
+    assert_eq!(app.stored_files().len(), 2);
+
+    let json = anon.show_crate("rate_limited");
+    assert_eq!(json.krate.max_version, "1.0.0");
+}
+
+#[test]
+fn publish_new_crate_override_loosens_ratelimit() {
+    let (app, anon, _, token) = TestApp::full()
+        // Most people get 1 new token every 1 day
+        .with_rate_limit(
+            LimitedAction::PublishNew,
+            Duration::from_secs(60 * 60 * 24),
+            1,
+        )
+        .with_database(TestDatabase::TestPool)
+        .with_token();
+
+    app.db(|conn| {
+        // Add an override so our user gets *2* new tokens (expires, y'know, sometime)
+        diesel::insert_into(publish_rate_overrides::table)
+            .values((
+                publish_rate_overrides::user_id.eq(token.as_model().user_id),
+                publish_rate_overrides::burst.eq(2),
+                publish_rate_overrides::expires_at.eq(None::<NaiveDateTime>),
+                publish_rate_overrides::action.eq(LimitedAction::PublishNew),
+            ))
+            .execute(conn)
+            .expect("Failed to add ratelimit override")
+    });
+
     let crate_to_publish = PublishBuilder::new("rate_limited1", "1.0.0");
     token.publish_crate(crate_to_publish).good();
 
     assert_eq!(app.stored_files().len(), 2);
 
-    // Uploading a second crate is limited
+    let json = anon.show_crate("rate_limited1");
+    assert_eq!(json.krate.max_version, "1.0.0");
+
     let crate_to_publish = PublishBuilder::new("rate_limited2", "1.0.0");
+    token.publish_crate(crate_to_publish).good();
+
+    assert_eq!(app.stored_files().len(), 4);
+
+    let json = anon.show_crate("rate_limited2");
+    assert_eq!(json.krate.max_version, "1.0.0");
+
+    let crate_to_publish = PublishBuilder::new("rate_limited3", "1.0.0");
     token
         .publish_crate(crate_to_publish)
         .assert_rate_limited(LimitedAction::PublishNew);
 
-    assert_eq!(app.stored_files().len(), 2);
+    assert_eq!(app.stored_files().len(), 4);
 
-    let response = anon.get::<()>("/api/v1/crates/rate_limited2");
+    let response = anon.get::<()>("/api/v1/crates/rate_limited3");
     assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
 
-    // Wait for the limit to be up
-    thread::sleep(Duration::from_millis(500));
+#[test]
+fn publish_new_crate_expired_override_ignored() {
+    let (app, anon, _, token) = TestApp::full()
+        // Most people get 1 new token every 1 day
+        .with_rate_limit(
+            LimitedAction::PublishNew,
+            Duration::from_secs(60 * 60 * 24),
+            1,
+        )
+        .with_database(TestDatabase::TestPool)
+        .with_token();
 
-    let crate_to_publish = PublishBuilder::new("rate_limited2", "1.0.0");
+    app.db(|conn| {
+        // Add an override so our user gets *2* new tokens (expires, y'know, sometime)
+        let just_now = Utc::now().naive_utc() - Duration::from_secs(1);
+        diesel::insert_into(publish_rate_overrides::table)
+            .values((
+                publish_rate_overrides::user_id.eq(token.as_model().user_id),
+                publish_rate_overrides::burst.eq(2),
+                publish_rate_overrides::expires_at.eq(just_now),
+                publish_rate_overrides::action.eq(LimitedAction::PublishNew),
+            ))
+            .execute(conn)
+            .expect("Failed to add ratelimit override")
+    });
+
+    let crate_to_publish = PublishBuilder::new("rate_limited1", "1.0.0");
     token.publish_crate(crate_to_publish).good();
 
-    let json = anon.show_crate("rate_limited2");
+    assert_eq!(app.stored_files().len(), 2);
+
+    let json = anon.show_crate("rate_limited1");
     assert_eq!(json.krate.max_version, "1.0.0");
 
-    assert_eq!(app.stored_files().len(), 4);
+    let crate_to_publish = PublishBuilder::new("rate_limited2", "1.0.0");
+    token
+        .publish_crate(crate_to_publish)
+        .assert_rate_limited(LimitedAction::PublishNew);
+
+    assert_eq!(app.stored_files().len(), 2);
+
+    let response = anon.get::<()>("/api/v1/crates/rate_limited2");
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
 }
 
 #[test]