WIP

src/handlers.rs

@@ -72,7 +72,7 @@ pub async fn store_file(base64_encoded_bytes: &str, pool: &storage:: DatabaseCon
     return Ok(warp::reply::json(&id).into_response());
 }
 
-pub async fn get_file(id: &str) -> Result<String, Rejection> {
+pub async fn get_file(id: &str) -> Result<String, Rejection> { // Doesn't return a response directly for testing purposes
     // Check that the ID is a valid UUID
     match Uuid::parse_str(id) {
         Ok(_) => (),

src/storage.rs

@@ -12,11 +12,11 @@ pub type DatabaseConnectionPool = r2d2::Pool<SqliteConnectionManager>;
 
 // Main
 
-const MAIN_TABLE: &str = "main";
+pub const MAIN_TABLE: &str = "main";
 
 lazy_static::lazy_static! {
 
-    static ref MAIN_POOL: DatabaseConnectionPool = {
+    pub static ref MAIN_POOL: DatabaseConnectionPool = {
         let file_name = format!("database.db");
         let db_manager = r2d2_sqlite::SqliteConnectionManager::file(file_name);
         return r2d2::Pool::new(db_manager).unwrap();
@@ -177,7 +177,7 @@ pub async fn prune_files_periodically() {
     let mut timer = tokio::time::interval(chrono::Duration::days(1).to_std().unwrap());
     loop {
         timer.tick().await;
-        tokio::spawn(async { prune_files().await; });
+        tokio::spawn(async { prune_files(FILE_EXPIRATION).await; });
     }
 }
 
@@ -243,16 +243,17 @@ async fn prune_pending_tokens() {
     println!("Pruned pending tokens.");
 }
 
-async fn prune_files() {
+pub async fn prune_files(file_expiration: i64) { // The expiration setting is passed in for testing purposes
     let rooms = match get_all_rooms().await {
         Ok(rooms) => rooms,
         Err(_) => return
     };
     for room in rooms {
         // It's not catastrophic if we fail to prune the database for a given room
+        println!("room name: {}", room);
         let pool = pool_by_room_name(&room);
         let now = chrono::Utc::now().timestamp();
-        let expiration = now - FILE_EXPIRATION;
+        let expiration = now - file_expiration;
         // Get a database connection and open a transaction
         let mut conn = match pool.get() {
             Ok(conn) => conn,
@@ -305,7 +306,7 @@ async fn get_all_rooms() -> Result<Vec<String>, Error> {
     // Get a database connection
     let conn = MAIN_POOL.get().map_err(|_| Error::DatabaseFailedInternally)?;
     // Query the database
-    let raw_query = format!("SELECT name FROM {}", MAIN_TABLE);
+    let raw_query = format!("SELECT id FROM {}", MAIN_TABLE);
     let mut query = conn.prepare(&raw_query).map_err(|_| Error::DatabaseFailedInternally)?;
     let rows = match query.query_map(params![], |row| {
         Ok(row.get(0)?)

src/tests.rs

@@ -17,7 +17,7 @@ fn perform_main_setup() {
     fs::create_dir_all("files").unwrap();
 }
 
-fn create_test_room() {
+fn set_up_test_room() {
     perform_main_setup();
     match fs::remove_file("rooms/test_room.db") {
         Ok(_) => (),
@@ -26,12 +26,18 @@ fn create_test_room() {
     let test_room = "test_room";
     storage::create_database_if_needed(test_room);
     fs::read("rooms/test_room.db").unwrap(); // Fail if this doesn't exist    
+    let pool: &storage::DatabaseConnectionPool = &storage::MAIN_POOL;
+    let mut conn = pool.get().unwrap();
+    let tx = conn.transaction().unwrap();
+    let stmt = format!("REPLACE INTO {} (id, name) VALUES (?1, ?2)", storage::MAIN_TABLE);
+    tx.execute(&stmt, params![ test_room, "Test Room" ]).unwrap();
+    tx.commit().unwrap();
 }
 
 #[test]
-fn test_file_storage_and_retrieval() {
+fn test_file_handling() {
     // Ensure the test room is set up
-    create_test_room();
+    set_up_test_room();
     // Test file storage
     let pool = storage::pool_by_room_name("test_room");
     aw!(handlers::store_file(TEST_FILE, &pool)).unwrap();
@@ -47,6 +53,20 @@ fn test_file_storage_and_retrieval() {
     // Retrieve the file and check the content
     let base64_encoded_file = aw!(handlers::get_file(id)).unwrap();
     assert_eq!(base64_encoded_file, TEST_FILE);
+    // Prune the file and check that it's gone
+    aw!(storage::prune_files(-60)); // Will evaluate to now + 60
+    match fs::read(format!("files/{}", id)) {
+        Ok(_) => assert!(false), // It should be gone now
+        Err(_) => ()
+    }
+    // Check that the file record is also gone
+    let mut conn = pool.get().unwrap();
+    let tx = conn.transaction().unwrap();
+    let raw_query = format!("SELECT id FROM {}", storage::FILES_TABLE);
+    let mut query = tx.prepare(&raw_query).unwrap();
+    let rows = query.query_map(params![], |row| { Ok(row.get(0)?) }).unwrap();
+    let ids: Vec<String> = rows.filter_map(|result| result.ok()).collect();
+    assert_eq!(ids.len(), 0);
 }
 
 // Data