Some tests for various API endpoints.

[wolnelektury.git] / src / catalogue / management / commands / savemedia.py

diff --git a/src/catalogue/management/commands/savemedia.py b/src/catalogue/management/commands/savemedia.py
index 216c5e3..ab4da51 100755 (executable)
--- a/src/catalogue/management/commands/savemedia.py
+++ b/src/catalogue/management/commands/savemedia.py
@@ -5,24 +5,22 @@
 import os.path
 
 from django.core.management.base import BaseCommand
+from django.db import transaction
 
 from catalogue.models import Book, BookMedia
 from catalogue.utils import ExistingFile
 
 
 class Command(BaseCommand):
-    help = "Saves uploaded media with a given book and a given name. If media has a source SHA1 info - matching media is replaced."
+    help = "Saves uploaded media with a given book and a given name. " \
+           "If media has a source SHA1 info - matching media is replaced."
     args = 'path slug name'
 
+    @transaction.atomic
     def handle(self, *args, **options):
-        from django.db import transaction
-
-        path, slug, name = args
-
-        # Start transaction management.
-        transaction.commit_unless_managed()
-        transaction.enter_transaction_management()
-        transaction.managed(True)
+        path, slug, name, part_name, index, parts_count = args
+        index = int(index)
+        parts_count = int(parts_count)
 
         book = Book.objects.get(slug=slug)
 
@@ -43,7 +41,7 @@ class Command(BaseCommand):
             bm = BookMedia(book=book, type=ext)
             print "Creating new media"
         bm.name = name
+        bm.part_name = part_name
+        bm.index = index
         bm.file.save(None, ExistingFile(path))
-        bm.save()
-        transaction.commit()
-        transaction.leave_transaction_management()
+        bm.save(parts_count=parts_count)