1 files changed, 18 insertions, 2 deletions
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index dcefcdab..c051ee0e 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -3,6 +3,8 @@ extern crate fatcat_api;
 extern crate chrono;
 #[macro_use]
 extern crate diesel;
+#[macro_use]
+extern crate diesel_migrations;
 extern crate dotenv;
 extern crate futures;
 extern crate uuid;
@@ -35,12 +37,15 @@ pub use errors::*;
 pub use self::errors::*;
 use diesel::pg::PgConnection;
 use diesel::prelude::*;
-use diesel::r2d2::ConnectionManager;
+use diesel::r2d2::{ConnectionManager, ManageConnection};
 use dotenv::dotenv;
 use iron::middleware::AfterMiddleware;
 use iron::{Request, Response};
 use std::env;
 
+#[cfg(feature = "postgres")]
+embed_migrations!("../migrations/");
+
 pub type ConnectionPool = diesel::r2d2::Pool<ConnectionManager<diesel::pg::PgConnection>>;
 
 /// Establish a direct database connection. Not currently used, but could be helpful for
@@ -55,8 +60,19 @@ pub fn establish_connection() -> PgConnection {
 /// Instantiate a new API server with a pooled database connection
 pub fn server() -> Result<api_server::Server> {
     dotenv().ok();
-    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+    let database_url = if cfg!(test) {
+        env::var("DATABASE_URL").expect("DATABASE_URL must be set")
+    } else {
+        env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set")
+    };
     let manager = ConnectionManager::<PgConnection>::new(database_url);
+    if cfg!(test) {
+        // run migrations; revert latest (dummy data); re-run latest
+        let conn = manager.connect().unwrap();
+        diesel_migrations::run_pending_migrations(&conn).unwrap();
+        diesel_migrations::revert_latest_migration(&conn).unwrap();
+        diesel_migrations::run_pending_migrations(&conn).unwrap();
+    }
     let pool = diesel::r2d2::Pool::builder()
         .build(manager)
         .expect("Failed to create database pool.");