Create test framework

1 files changed, 66 insertions, 58 deletions

diff --git a/models/artist.py b/models/artist.py
index fee081c..a76b2ee 100644
--- a/models/artist.py
+++ b/models/artist.py
@@ -1,44 +1,50 @@
 from common import utils
 from db.db_manager import DbManager
+from models.base import BaseModel
 
 
-class Artist:
-    def __init__(self, id=None, **kwargs):
+class Artist(BaseModel):
+
+    def __init__(self, id=None, db=None, **kwargs):
+        if db:
+            self.db = db
+
         if id is not None:
-            db = DbManager()
-            cursor = db.cursor()
-
-            for row in cursor.execute("SELECT * FROM artist WHERE id = ?",
-                                      (id,)):
-                setattr(self, "id", id)
-                setattr(self, "name", row[1])
-                setattr(self, "sortname", row[2])
-                setattr(self, "musicbrainz_artistid", row[3])
+            for row in self.db.execute("SELECT * FROM artist WHERE id = ?",
+                                       (id,)):
+                for key in ["id", "name", "sortname", "musicbrainz_artistid"]:
+                    setattr(self, key, row[key])
         else:
             for (key, value) in kwargs.items():
                 setattr(self, key, value)
 
     def delete(self):
-        db = DbManager()
-        cursor = db.cursor()
-
         for album in self.albums:
             album.delete()
 
-        cursor.execute("BEGIN TRANSACTION")
+        with self.db.conn:
+            delete_artist = "DELETE FROM artist WHERE id = ?"
+            self.db.execute(delete_artist, (self.id,))
 
-        delete_sql = "DELETE FROM artist WHERE id = ?"
-        cursor.execute(delete_sql, (self.id,))
+            delete_track_rel = "DELETE FROM artist_track WHERE artist_id = ?"
+            self.db.execute(delete_track_rel, (self.id,))
 
-        delete_track_rel_sql = "DELETE FROM artist_track WHERE artist_id = ?"
-        cursor.execute(delete_track_rel_sql, (self.id,))
+            delete_album_rel = "DELETE FROM album_artist WHERE artist_id = ?"
+            self.db.execute(delete_album_rel, (self.id,))
 
-        delete_album_rel_sql = "DELETE FROM album_artist WHERE artist_id = ?"
-        cursor.execute(delete_album_rel_sql, (self.id,))
+        return True
 
-        cursor.execute("COMMIT TRANSACTION")
+    @property
+    def db(self):
+        try:
+            return self._db
+        except AttributeError:
+            self._db = DbManager()
+            return self._db
 
-        return True
+    @db.setter
+    def db(self, db):
+        self._db = db
 
     @property
     def tracks(self):
@@ -47,16 +53,16 @@ class Artist:
         if not hasattr(self, "_tracks"):
             setattr(self, "_tracks", [])
 
-            db = DbManager()
-            cursor = db.cursor()
-
-            for row in cursor.execute("""SELECT track.* FROM track
-                                      INNER JOIN artist_track ON track.id =
-                                      artist_track.track_id WHERE artist_id = ?
-                                      ORDER BY name ASC""", (self.id,)):
+            for row in self.db.execute("SELECT track.* FROM track INNER "
+                                       "JOIN artist_track ON track.id = "
+                                       "artist_track.track_id WHERE "
+                                       "artist_id = ? ORDER BY name ASC",
+                                       (self.id,)):
 
-                track = Track(id=row[0], tracknumber=row[1], name=row[2],
-                              grouping=row[3], filename=row[4])
+                track = Track(id=row["id"], db=self.db,
+                              tracknumber=row["tracknumber"], name=row["name"],
+                              grouping=row["grouping"],
+                              filename=row["filename"])
                 self._tracks.append(track)
 
         return self._tracks
@@ -68,14 +74,13 @@ class Artist:
         if not hasattr(self, "_albums"):
             setattr(self, "_albums", [])
 
-            db = DbManager()
-            cursor = db.cursor()
-
-            for row in cursor.execute("""SELECT album.* FROM album
-                                      INNER JOIN album_artist ON album.id =
-                                      album_artist.album_id WHERE artist_id = ?
-                                      ORDER BY date ASC""", (self.id,)):
-                album = Album(id=row[0], name=row[1], date=row[2])
+            for row in self.db.execute("SELECT album.* FROM album INNER "
+                                       "JOIN album_artist ON album.id = "
+                                       "album_artist.album_id WHERE "
+                                       "artist_id = ? ORDER BY date ASC",
+                                       (self.id,)):
+                album = Album(id=row["id"], db=self.db, name=row["name"],
+                              date=row["date"])
                 self._albums.append(album)
 
         return self._albums
@@ -89,17 +94,16 @@ class Artist:
                 dirty_attributes[attr] = value
 
         if len(dirty_attributes) > 0:
-            db = DbManager()
-            cursor = db.cursor()
-
             set_clause = utils.update_clause_from_dict(dirty_attributes)
 
             dirty_attributes[id] = self.id
 
             sql = " ".join(("UPDATE artist"), set_clause, "WHERE id = :id")
-            cursor.execute(sql, dirty_attributes)
 
-    def search(**search_params):
+            with self.db.conn:
+                self.db.execute(sql, dirty_attributes)
+
+    def search(db=None, **search_params):
         """Find an artist with the given params
 
         Args:
@@ -109,8 +113,8 @@ class Artist:
         """
         artists = []
 
-        db = DbManager()
-        cursor = db.cursor()
+        if not db:
+            db = DbManager()
 
         # unpack search params
         where_params = {}
@@ -124,21 +128,24 @@ class Artist:
         result = []
         if where_clause:
             statement = " ".join(("SELECT * FROM artist", where_clause))
-            result = cursor.execute(statement, value_params)
+            result = db.execute(statement, value_params)
         else:
-            result = cursor.execute("SELECT * FROM artist")
+            result = db.execute("SELECT * FROM artist")
 
         for row in result:
             artists.append(
-                Artist(id=row[0], name=row[1], sortname=row[2],
-                       musicbrainz_artistid=row[3])
+                Artist(id=row["id"], db=db, name=row["name"],
+                       sortname=row["sortname"],
+                       musicbrainz_artistid=row["musicbrainz_artistid"])
             )
 
         return artists
 
-    def all(order="sortname", direction="ASC", limit=None, offset=None):
-        db = DbManager()
-        cursor = db.cursor()
+    def all(db=None, order="sortname", direction="ASC", limit=None,
+            offset=None):
+        if not db:
+            db = DbManager()
+
         artists = []
 
         select_string = "SELECT * FROM artist ORDER BY %s %s" % (order,
@@ -146,14 +153,15 @@ class Artist:
 
         if limit is not None and offset is not None:
             select_string = " ".join((select_string,
-                                     "LIMIT %s OFFSET %s" % (limit, offset)))
+                                      "LIMIT %s OFFSET %s" % (limit, offset)))
 
-        result = cursor.execute(select_string)
+        result = db.execute(select_string)
 
         for row in result:
             artists.append(
-                Artist(id=row[0], name=row[1], sortname=row[2],
-                       musicbrainz_artistid=row[3])
+                Artist(id=row["id"], db=db, name=row["name"],
+                       sortname=row["sortname"],
+                       musicbrainz_artistid=row["musicbrainz_artistid"])
             )
 
         return artists