Merge pull request #232 from mkurek/feature/service-san-plugin

san collect plugin

src/ralph_scrooge/models/usage.py

@@ -263,7 +263,7 @@ class DailyUsage(db.Model):
         blank=False,
     )
     value = db.FloatField(verbose_name=_("value"), default=0)
-    type = db.ForeignKey(UsageType, verbose_name=_("type"))
+    type = db.ForeignKey(UsageType, verbose_name=_("daily_usages"))
     warehouse = db.ForeignKey('Warehouse', null=True, on_delete=db.PROTECT)
     remarks = db.TextField(
         verbose_name=_("Remarks"),
@@ -279,7 +279,7 @@ class Meta:
 
     def __unicode__(self):
         return '{0}/{1} ({2}) {3}'.format(
-            self.pricing_object,
+            self.daily_pricing_object,
             self.type,
             self.date,
             self.value,

src/ralph_scrooge/plugins/collect/san.py

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import logging
+
+from ralph.util import plugin
+from ralph.util.api_pricing import get_fc_cards
+from ralph_scrooge.models import (
+    AssetInfo,
+    DailyAssetInfo,
+    DailyUsage,
+    UsageType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AssetInfoNotFound(Exception):
+    pass
+
+
+class DailyAssetInfoNotFound(Exception):
+    pass
+
+
+def update_usage(daily_asset_info, service, date, value, usage_type):
+    """
+    Saves single record to model
+    """
+    usage, created = DailyUsage.objects.get_or_create(
+        date=date,
+        type=usage_type,
+        daily_pricing_object=daily_asset_info,
+        defaults=dict(
+            service=service,
+        )
+    )
+    usage.service = service
+    usage.value = value
+    usage.save()
+    return created
+
+
+def update_san(data, date, usage_type):
+    """
+    Updates single SAN usage type record
+    """
+    try:
+        asset_info = AssetInfo.objects.get(device_id=data['device_id'])
+        daily_asset_info = asset_info.dailyassetinfo_set.get(date=date)
+        service = daily_asset_info.service
+        return update_usage(
+            daily_asset_info,
+            service,
+            date,
+            1,
+            usage_type,
+        )
+    except AssetInfo.DoesNotExist:
+        raise AssetInfoNotFound()
+    except DailyAssetInfo.DoesNotExist:
+        raise DailyAssetInfoNotFound()
+
+
+def get_usage_type():
+    """
+    Returns SAN usage type
+    """
+    return UsageType.objects.get_or_create(
+        symbol='san',
+        defaults=dict(
+            name='SAN',
+        )
+    )[0]
+
+
+@plugin.register(chain='scrooge', requires=['asset', 'service'])
+def san(today, **kwargs):
+    """
+    Updates SAN usages from Ralph
+    """
+    usage_type = get_usage_type()
+    new_san = updated = total = 0
+    for data in get_fc_cards():
+        try:
+            if update_san(data, today, usage_type):
+                new_san += 1
+            else:
+                updated += 1
+        except AssetInfoNotFound:
+            logger.warning('Device {} not found'.format(data['device_id']))
+        except DailyAssetInfoNotFound:
+            logger.warning(
+                'DailyAssetInfo for id {} and date {} not found'.format(
+                    data['device_id'],
+                    today,
+                )
+            )
+        total += 1
+    return (
+        True,
+        '{} new SAN usages, {} updated, {} total'.format(
+            new_san,
+            updated,
+            total,
+        )
+    )

src/ralph_scrooge/tests/plugins/collect/samples/san.py

@@ -0,0 +1,18 @@
+SAMPLE_SAN = [
+    {
+        'id': 1,
+        'device_id': 10,
+    },
+    {
+        'id': 2,
+        'device_id': 20
+    },
+    {
+        'id': 3,
+        'device_id': 30
+    },
+    {
+        'id': 4,
+        'device_id': 40
+    },
+]

src/ralph_scrooge/tests/plugins/collect/test_san.py

@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import datetime
+import mock
+
+from django.test import TestCase
+
+from ralph_scrooge.plugins.collect.san import (
+    AssetInfoNotFound,
+    DailyAssetInfoNotFound,
+    san as san_plugin,
+    get_usage_type,
+    update_san,
+    update_usage,
+)
+from ralph_scrooge.tests.utils.factory import DailyAssetInfoFactory
+from ralph_scrooge.tests.plugins.collect.samples.san import SAMPLE_SAN
+
+
+class TestSanCollectPlugin(TestCase):
+    def setUp(self):
+        self.today = datetime.date(2014, 7, 1)
+        self.san_usage_type = get_usage_type()
+
+    def test_update_usage(self):
+        daily_asset_info = DailyAssetInfoFactory(date=self.today)
+        self.assertTrue(update_usage(
+            daily_asset_info,
+            daily_asset_info.service,
+            self.today,
+            1,
+            self.san_usage_type,
+        ))
+        self.assertEquals(self.san_usage_type.dailyusage_set.count(), 1)
+        daily_usage = self.san_usage_type.dailyusage_set.all()[:1].get()
+        self.assertEquals(daily_usage.date.date(), self.today)
+        self.assertEquals(daily_usage.service, daily_asset_info.service)
+        self.assertEquals(daily_usage.type, self.san_usage_type)
+        self.assertEquals(daily_usage.value, 1)
+
+        # update
+        self.assertFalse(update_usage(
+            daily_asset_info,
+            daily_asset_info.service,
+            self.today,
+            2,
+            self.san_usage_type,
+        ))
+        self.assertEquals(self.san_usage_type.dailyusage_set.count(), 1)
+        daily_usage = self.san_usage_type.dailyusage_set.all()[:1].get()
+        self.assertEquals(daily_usage.value, 2)
+
+    @mock.patch('ralph_scrooge.plugins.collect.san.update_usage')
+    def test_update_san(self, update_usage_mock):
+        update_usage_mock.return_value = True
+        daily_asset_info = DailyAssetInfoFactory(date=self.today)
+        self.assertTrue(update_san(
+            {'id': 1, 'device_id': daily_asset_info.asset_info.device_id},
+            self.today,
+            self.san_usage_type
+        ))
+        update_usage_mock.assert_called_with(
+            daily_asset_info,
+            daily_asset_info.service,
+            self.today,
+            1,
+            self.san_usage_type,
+        )
+
+    def test_update_san_no_daily_asset_info(self):
+        daily_asset_info = DailyAssetInfoFactory(date=self.today)
+        with self.assertRaises(DailyAssetInfoNotFound):
+            update_san(
+                {'id': 1, 'device_id': daily_asset_info.asset_info.device_id},
+                self.today + datetime.timedelta(days=1),
+                self.san_usage_type
+            )
+
+    def test_update_san_no_asset_info(self):
+        daily_asset_info = DailyAssetInfoFactory.build(date=self.today)
+        with self.assertRaises(AssetInfoNotFound):
+            update_san(
+                {'id': 1, 'device_id': daily_asset_info.asset_info.device_id},
+                self.today,
+                self.san_usage_type
+            )
+
+    @mock.patch('ralph_scrooge.plugins.collect.san.update_san')
+    @mock.patch('ralph_scrooge.plugins.collect.san.get_fc_cards')
+    def test_batch_update(self, get_fc_cards_mock, update_san_mock):
+        def sample_get_fc_cards():
+            return SAMPLE_SAN
+
+        def sample_update_san_mock(data, date, usage_type):
+            responses = {
+                20: False,
+                30: AssetInfoNotFound(),
+                40: DailyAssetInfoNotFound(),
+            }
+            response = responses.get(data['device_id'], True)
+            if isinstance(response, Exception):
+                raise response
+            return response
+
+        get_fc_cards_mock.side_effect = sample_get_fc_cards
+        update_san_mock.side_effect = sample_update_san_mock
+        self.assertEquals(
+            san_plugin(today=self.today),
+            (True, '1 new SAN usages, 1 updated, 4 total')
+        )

src/ralph_scrooge/tests/utils/factory.py

@@ -6,8 +6,10 @@
 from __future__ import unicode_literals
 
 import datetime
+import random
 
 from factory import (
+    fuzzy,
     lazy_attribute,
     Sequence,
     SubFactory,
@@ -16,6 +18,9 @@
 
 from ralph_scrooge import models
 
+MIN_FACTORY_DATE = datetime.date(2014, 7, 1)
+MAX_FACTORY_DATE = datetime.date(2014, 7, 31)
+
 
 class WarehouseFactory(DjangoModelFactory):
     FACTORY_FOR = models.Warehouse
@@ -58,9 +63,20 @@ class AssetInfoFactory(PricingObjectFactory):
     sn = Sequence(lambda n: n)
     barcode = Sequence(lambda n: n)
     asset_id = Sequence(lambda n: n)
+    device_id = Sequence(lambda n: n)
     warehouse = SubFactory(WarehouseFactory)
 
 
+class DailyAssetInfoFactory(DailyPricingObjectFactory):
+    FACTORY_FOR = models.DailyAssetInfo
+
+    asset_info = SubFactory(AssetInfoFactory)
+    depreciation_rate = fuzzy.FuzzyDecimal(0, 50)
+    is_depreciated = fuzzy.FuzzyAttribute(lambda: random.random() < 0.5)
+    price = fuzzy.FuzzyDecimal(0, 1000)
+    date = fuzzy.FuzzyDate(MIN_FACTORY_DATE, MAX_FACTORY_DATE)
+
+
 class OwnerFactory(DjangoModelFactory):
     FACTORY_FOR = models.Owner