gistfile1.txt

diff --git a/Cargo.toml b/Cargo.toml
index ddcf743..5ca90f2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,9 +11,10 @@ reqwest = "0.9.19"
 serde = "1.0.98"
 serde_json = "1.0.40"
 regex = "1.2.0"
-diesel = {version = "1.4.2", features = ["sqlite"]}
+diesel = {version = "1.4.2", features = ["sqlite", "r2d2"]}
 diesel_migrations = "1.4.0"
 rocket = "0.4.2"
 rocket_contrib = "0.4.2"
 dotenv = "0.14"
 indicatif = "0.11.0"
+r2d2 = "0.8.5"
diff --git a/src/db/mod.rs b/src/db/mod.rs
index 0e55b43..2170380 100644
--- a/src/db/mod.rs
+++ b/src/db/mod.rs
@@ -5,11 +5,18 @@ mod scryfall;
 mod sql;
 pub mod sqlite;
 
-use diesel::{prelude::*, result::Error as DieselError, sqlite::SqliteConnection};
+use diesel::{
+    prelude::*,
+    result::Error as DieselError,
+    sqlite::SqliteConnection,
+    sql_query,
+    r2d2::ConnectionManager
+};
+
 use diesel_migrations::RunMigrationsError;
 use reqwest::Error as ReqwestError;
 
-use std::{collections::HashMap, convert::TryFrom};
+use std::{collections::HashMap, convert::TryFrom, env, ops::Deref};
 
 pub use self::queries::QueryError;
 use crate::model::{CardData, DeckCard, ParseError};
@@ -25,6 +32,7 @@ pub enum DBError {
     Request(u16),
     SQLiteConnection(ConnectionError),
     DieselMigrations(RunMigrationsError),
+    Pool,
 }
 
 impl DBError {
@@ -80,5 +88,30 @@ impl TryFrom<(CardDataList, &SqliteConnection)> for CardDataList {
 }
 
 pub struct CardDB {
-    connection: SqliteConnection,
+    connection: r2d2::PooledConnection<ConnectionManager<SqliteConnection>>,
+}
+
+impl CardDB {
+    pub fn new() -> Result<Self, DBError> {
+        let db_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+        let connection =
+            SqliteConnection::establish(&*db_url).map_err(DBError::SQLiteConnection)?;
+        // enable foreign keys
+        Self::from_connection(connection)
+    }
+
+    pub fn from_connection(connection: SqliteConnection) -> Result<Self, DBError> {
+        sql_query("PRAGMA foreign_keys = ON;")
+            .execute(&connection)
+            .map_err(DBError::Bug)?;
+        Ok(CardDB { connection })
+    }
+}
+
+impl Deref for CardDB {
+    type Target = SqliteConnection;
+
+    fn deref(&self) -> &Self::Target {
+        &self.connection
+    }
 }
diff --git a/src/db/queries.rs b/src/db/queries.rs
index 75ba8fa..6d7c16b 100644
--- a/src/db/queries.rs
+++ b/src/db/queries.rs
@@ -1,11 +1,11 @@
-use std::{convert::TryFrom, env};
+use std::convert::TryFrom;
 
 use diesel::{
     dsl::{exists, sql},
     insert_into,
     prelude::*,
     result::{DatabaseErrorKind, Error as DieselError},
-    select, sql_query,
+    select,
     sql_types::{BigInt, Bool},
     update,
 };
@@ -55,15 +55,10 @@ macro_rules! check_exists_id {
 }
 
 impl CardDB {
-    pub fn new() -> Result<Self, DBError> {
-        let db_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
-        let connection =
-            SqliteConnection::establish(&*db_url).map_err(DBError::SQLiteConnection)?;
-        // enable foreign keys
-        sql_query("PRAGMA foreign_keys = ON;")
-            .execute(&connection)
-            .map_err(DBError::Bug)?;
-        Ok(Self { connection })
+    pub fn setup_database(&self) -> Result<(), DBError> {
+        setup_database(&self.connection).map_err(DBError::Bug)?;
+        run_pending_migrations(&self.connection).map_err(DBError::DieselMigrations)?;
+        Ok(())
     }
 
     pub fn initialize<F>(&self, callback: Option<F>) -> Result<(), DBError>
@@ -85,8 +80,7 @@ impl CardDB {
             Ok(())
         } else {
             println!("Initializing DB from scratch...");
-            setup_database(&self.connection).map_err(DBError::Bug)?;
-            run_pending_migrations(&self.connection).map_err(DBError::DieselMigrations)?;
+            self.setup_database()?;
             self.initialize_cards(callback)?;
             Ok(())
         }
@@ -240,13 +234,14 @@ impl CardDB {
 
     pub fn add_to_collection(&self, _version_id: i64, _number: i32) -> Result<(), QueryError> {
         use crate::db::schema::collection_cards::dsl::*;
-        let res = update(collection_cards)
+        let updated = update(collection_cards)
             .filter(version_id.eq(_version_id))
             .set(number.eq(number + _number))
-            .execute(&self.connection);
+            .execute(&self.connection)
+            .map_err(QueryError::from_diesel)?;
 
         // Error means that there isn't such an entry yet
-        if res.is_err() {
+        if updated == 0 {
             if _number <= 0 {
                 return Err(QueryError::InvalidParameter);
             }
@@ -256,7 +251,12 @@ impl CardDB {
                     number: _number,
                 })
                 .execute(&self.connection)
-                .map_err(QueryError::from_diesel)?;
+                .map_err(|e| match e {
+                    DieselError::DatabaseError(DatabaseErrorKind::ForeignKeyViolation, _) => {
+                        QueryError::InvalidParameter
+                    }
+                    other => QueryError::from_diesel(other),
+                })?;
         }
         Ok(())
     }
diff --git a/src/web/collection.rs b/src/web/collection.rs
new file mode 100644
index 0000000..e89ac97
--- /dev/null
+++ b/src/web/collection.rs
@@ -0,0 +1,25 @@
+use rocket_contrib::json::{Json, JsonError};
+
+use crate::db::CardDB;
+
+use super::APIError;
+
+#[derive(Deserialize)]
+pub struct AddCollectionData {
+    id: i64,
+    number: i32,
+}
+
+#[post("/collection", data = "<json>")]
+pub fn add_to_collection(
+    json: Result<Json<AddCollectionData>, JsonError>,
+    db: CardDB,
+) -> Result<(), APIError> {
+    match json {
+        Err(_) => Err(APIError::Params(String::from("Need a valid id and number"))),
+        Ok(data) => {
+            db.add_to_collection(data.id, data.number)
+                .map_err(APIError::Query)
+        }
+    }
+}
diff --git a/src/web/decks.rs b/src/web/decks.rs
index ceb3cbb..313e428 100644
--- a/src/web/decks.rs
+++ b/src/web/decks.rs
@@ -1,8 +1,9 @@
+
 use rocket_contrib::json::{Json, JsonError};
 
 use crate::db::{CardDB, CardDataList, QueryError};
 use crate::importers::mtggoldfish;
-use crate::model::{Deck, Format, NewDeck};
+use crate::model::{Format, Deck, NewDeck};
 
 use super::APIError;
 
@@ -27,6 +28,7 @@ pub struct CardDeckResult {
     sideboard: CardDataList,
 }
 
+
 #[post("/deck/import/<provider>/<deck_id>")]
 pub fn import_deck(
     provider: String,
diff --git a/src/web/mod.rs b/src/web/mod.rs
index f08a2a0..9feded1 100644
--- a/src/web/mod.rs
+++ b/src/web/mod.rs
@@ -1,16 +1,25 @@
+mod collection;
 mod decks;
+#[cfg(test)]
+mod tests;
 
+use diesel::{
+    r2d2::ConnectionManager,
+    sqlite::SqliteConnection
+};
+
+use r2d2;
 use rocket::{
     http::ContentType,
     http::Status,
     request::{FromRequest, Outcome},
     response,
     response::Responder,
-    Request, Response,
+    Request, Response, Rocket, State
 };
 use serde::{ser::SerializeStruct, Serialize, Serializer};
 use serde_json::to_string_pretty;
-use std::io::Cursor;
+use std::{env, io::Cursor};
 
 use crate::db::{CardDB, DBError, QueryError};
 use crate::importers::Error as ImporterError;
@@ -26,6 +35,13 @@ pub enum APIError {
     Importer(ImporterError),
 }
 
+type Pool = r2d2::Pool<ConnectionManager<SqliteConnection>>;
+
+pub fn init_pool() -> Pool {
+    let manager = ConnectionManager::<SqliteConnection>::new(env::var("DATABASE_URL").expect("DATABASE_URL env var"));
+    Pool::new(manager).expect("db pool")
+}
+
 impl Serialize for APIError {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -76,13 +92,14 @@ impl<'r> Responder<'r> for APIError {
     }
 }
 
-impl<'a, 'r> FromRequest<'a, 'r> for CardDB {
+impl<'a, 'r, 't> FromRequest<'a, 'r> for CardDB {
     type Error = DBError;
 
-    fn from_request(_: &'a Request<'r>) -> Outcome<Self, Self::Error> {
-        match CardDB::new() {
-            Ok(d) => Outcome::Success(d),
-            Err(e) => Outcome::Failure((Status::InternalServerError, e)),
+    fn from_request(request: &'a Request<'r>) -> Outcome<Self, Self::Error> {
+        let pool = request.guard::<State<Pool>>().map_err(DBError::Pool)?;
+        match pool.get() {
+            Ok(conn) => Outcome::Success(CardDB::from_connection(conn).map_err(DBError::Pool)?),
+            Err(e) => Outcome::Failure((Status::ServiceUnavailable, DBError::SQLiteConnection(e)),
         }
     }
 }
@@ -92,7 +109,7 @@ fn index<'t>() -> &'t str {
     "Hello world again!"
 }
 
-pub fn run_server() {
+fn rocket() -> Rocket {
     rocket::ignite()
         .mount(
             "/",
@@ -102,8 +119,12 @@ pub fn run_server() {
                 decks::create_deck,
                 decks::get_deck,
                 decks::add_to_deck,
-                decks::remove_from_deck
+                decks::remove_from_deck,
+                collection::add_to_collection
             ],
         )
-        .launch();
+}
+
+pub fn run_server() {
+    rocket().launch();
 }
diff --git a/src/web/tests.rs b/src/web/tests.rs
new file mode 100644
index 0000000..8165ade
--- /dev/null
+++ b/src/web/tests.rs
@@ -0,0 +1,38 @@
+use std::env;
+
+use super::rocket;
+use crate::db::CardDB;
+use rocket::local::Client;
+use rocket::http::Status;
+
+macro_rules! setup {
+    ($client: ident, $code: block) => {
+        env::set_var("DATABASE_URL", ":memory:");
+        let $client = Client::new(rocket()).expect("valid rocket instance");
+        let db = CardDB::new().expect("Database connection");
+
+        db.setup_database().expect("Database setup");
+        $code
+    };
+}
+
+#[test]
+fn create_deck_fail() {
+    setup! (client, {
+        let response = client.post("/deck").dispatch();
+        assert_eq!(response.status(), Status::BadRequest);
+
+        let mut response = client.post("/deck").body(r#"{"name": "Burn", "format": "Foo"}"#).dispatch();
+        assert!(response.body_string().unwrap().contains("Need a valid format and name"));
+    });
+}
+
+#[test]
+fn create_deck_ok() {
+    setup! (client, {
+        let mut response = client.post("/deck").body(r#"{"name": "Burn", "format": "Standard"}"#).dispatch();
+        println!("{:?}", response.body().unwrap().into_string().unwrap());
+        assert_eq!(response.status(), Status::Ok);
+    });
+}
+