generic process addon task command, flag binary task (bug 655115)

apps/addons/management/commands/process_addons.py

@@ -0,0 +1,35 @@
+from optparse import make_option
+
+from django.core.management.base import BaseCommand
+
+from addons.models import Addon
+from amo.utils import chunked
+from devhub.tasks import flag_binary
+
+from celery.messaging import establish_connection
+
+tasks = {
+    'flag_binary': flag_binary,
+}
+
+
+class Command(BaseCommand):
+    """
+    A generic command to run a task on *all* addons.
+    Add tasks to the reg dictionary.
+    """
+    help = 'Corrects the binary status of an addon.'
+
+    option_list = BaseCommand.option_list + (
+        make_option('--task', action='store', type='string',
+                    dest='task', help='Run task on all addons.'),
+    )
+
+    def handle(self, *args, **options):
+        pks = Addon.objects.values_list('pk', flat=True).order_by('id')
+        task = tasks.get(options.get('task'))
+        if not task:
+            raise ValueError('Unknown task: %s' % ', '.join(tasks.keys()))
+        with establish_connection():
+            for chunk in chunked(pks, 100):
+                task.delay(chunk)

apps/devhub/tasks.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import sys
@@ -7,6 +8,7 @@ from django.conf import settings
 from django.core.management import call_command
 from celeryutils import task
 
+from addons.models import Addon
 import amo
 from amo.decorators import write, set_modified_on
 from amo.utils import resize_image
@@ -90,6 +92,28 @@ def run_validator(file_path, for_appversions=None):
                     spidermonkey=settings.SPIDERMONKEY)
 
 
+@task(rate_limit='1/m')
+@write
+def flag_binary(ids, **kw):
+    log.info('[%s@%s] Flagging binary addons starting with id: %s...'
+             % (len(ids), flag_binary.rate_limit, ids[0]))
+    addons = Addon.objects.filter(pk__in=ids).no_transforms()
+
+    for addon in addons:
+        try:
+            log.info('Validating addon with id: %s' % addon.pk)
+            file = File.objects.filter(version__addon=addon).latest('created')
+            result = run_validator(file.file_path)
+            binary = (json.loads(result)['metadata']
+                          .get('contains_binary_extension', False))
+            log.info('Setting binary for addon with id: %s to %s'
+                     % (addon.pk, binary))
+            addon.update(binary=binary)
+        except Exception, err:
+            log.error('Failed to run validation on addon id: %s, %s'
+                      % (addon.pk, err))
+
+
 @task
 @set_modified_on
 def resize_icon(src, dst, size, **kw):

apps/devhub/tests/test_tasks.py

@@ -10,8 +10,9 @@ import test_utils
 from nose.tools import eq_
 from PIL import Image
 
+from addons.models import Addon
 from files.models import FileUpload
-from devhub.tasks import resize_icon, validator
+from devhub.tasks import flag_binary, resize_icon, validator
 
 
 def test_resize_icon_shrink():
@@ -112,3 +113,29 @@ class TestValidator(test_utils.TestCase):
             validator(self.upload.pk)
         error = self.get_upload().task_error
         assert error.startswith('Traceback (most recent call last)'), error
+
+
+class TestFlagBinary(test_utils.TestCase):
+    fixtures = ['base/addon_3615']
+
+    def setUp(self):
+        self.addon = Addon.objects.get(pk=3615)
+        self.addon.update(binary=False)
+
+    @mock.patch('devhub.tasks.run_validator')
+    def test_flag_binary(self, _mock):
+        _mock.return_value = '{"metadata":{"contains_binary_extension": 1}}'
+        flag_binary([self.addon.pk])
+        eq_(Addon.objects.get(pk=self.addon.pk).binary, True)
+
+    @mock.patch('devhub.tasks.run_validator')
+    def test_flag_not_binary(self, _mock):
+        _mock.return_value = '{"metadata":{"contains_binary_extension": 0}}'
+        flag_binary([self.addon.pk])
+        eq_(Addon.objects.get(pk=self.addon.pk).binary, False)
+
+    @mock.patch('devhub.tasks.run_validator')
+    def test_flag_error(self, _mock):
+        _mock.side_effect = RuntimeError()
+        flag_binary([self.addon.pk])
+        eq_(Addon.objects.get(pk=self.addon.pk).binary, False)