assets on raise xD

allegro · Aug 5, 2014 · a025b7b · a025b7b
1 parent 61ba085
commit a025b7b
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 38 deletions.
diff --git a/src/ralph_scrooge/plugins/collect/asset.py b/src/ralph_scrooge/plugins/collect/asset.py
@@ -27,6 +27,20 @@
 logger = logging.getLogger(__name__)
 
 
+class ServiceDoesNotExistError(Exception):
+    """
+    Raise this exception when service does not exist
+    """
+    pass
+
+
+class WarehouseDoesNotExistError(Exception):
+    """
+    Raise this exception when warehouse does not exist
+    """
+    pass
+
+
 def create_pricing_object(service, data):
     """
     Create pricing object
@@ -160,18 +174,12 @@ def update_assets(data, date, usages):
     try:
         service = Service.objects.get(ci_uid=data['service_ci_uid'])
     except Service.DoesNotExist:
-        logger.error('Service {0} does not exist'.format(
-            data['service_ci_uid'],
-        ))
-        return (False, False)
+        raise ServiceDoesNotExistError()
 
     try:
         warehouse = Warehouse.objects.get(id_from_assets=data['warehouse_id'])
     except Warehouse.DoesNotExist:
-        logger.error('Warehouse {0} does not exist'.format(
-            data['warehouse_id']
-        ))
-        return (False, False)
+        raise WarehouseDoesNotExistError()
 
     asset_info, pricing_object, new_created = get_asset_and_pricing_object(
         service,
@@ -214,7 +222,7 @@ def update_assets(data, date, usages):
         date,
         warehouse,
     )
-    return (True, new_created)
+    return new_created
 
 
 def get_usage(symbol, name, by_warehouse, by_cost, average):
@@ -276,12 +284,22 @@ def asset(**kwargs):
 
     new = update = total = 0
     for data in get_assets(date):
-        result = update_assets(data, date, usages)
-        if result[0]:
-            if result[1]:
-                new += 1
-            else:
-                update += 1
         total += 1
+        try:
+            result = update_assets(data, date, usages)
+        except ServiceDoesNotExistError:
+            logger.error('Service {0} does not exist'.format(
+                data['service_ci_uid'],
+            ))
+            continue
+        except WarehouseDoesNotExistError:
+            logger.error('Warehouse {0} does not exist'.format(
+                data['warehouse_id']
+            ))
+            continue
+        if result:
+            new += 1
+        else:
+            update += 1
 
     return True, '{0} new, {1} updated, {2} total'.format(new, update, total)
diff --git a/src/ralph_scrooge/tests/plugins/collect/test_asset.py b/src/ralph_scrooge/tests/plugins/collect/test_asset.py
@@ -136,32 +136,30 @@ def test_update_usage(self):
 
     def test_update_assets_when_service_does_not_exist(self):
         self.data['service_ci_uid'] = 2
-        self.assertEqual(
-            asset.update_assets(
-                self.data,
-                self.date,
-                {
-                    'core': UsageTypeFactory.create(),
-                    'power_consumption': UsageTypeFactory.create(),
-                    'collocation': UsageTypeFactory.create(),
-                }
-            ),
-            (False, False)
+        self.assertRaises(
+            asset.ServiceDoesNotExistError,
+            asset.update_assets,
+            data=self.data,
+            date=self.date,
+            usages={
+                'core': UsageTypeFactory.create(),
+                'power_consumption': UsageTypeFactory.create(),
+                'collocation': UsageTypeFactory.create(),
+            },
         )
 
     def test_update_assets_when_warehouse_does_not_exist(self):
         self.data['warehouse_id'] = 2
-        self.assertEqual(
-            asset.update_assets(
-                self.data,
-                self.date,
-                {
-                    'core': UsageTypeFactory.create(),
-                    'power_consumption': UsageTypeFactory.create(),
-                    'collocation': UsageTypeFactory.create(),
-                }
-            ),
-            (False, False)
+        self.assertRaises(
+            asset.WarehouseDoesNotExistError,
+            asset.update_assets,
+            data=self.data,
+            date=self.date,
+            usages={
+                'core': UsageTypeFactory.create(),
+                'power_consumption': UsageTypeFactory.create(),
+                'collocation': UsageTypeFactory.create(),
+            },
         )
 
     def test_update_assets(self):
@@ -175,7 +173,7 @@ def test_update_assets(self):
                     'collocation': UsageTypeFactory.create(),
                 }
             ),
-            (True, True)
+            True
         )
         self.assertEqual(
             DailyUsage.objects.all().count(),