diff --git a/koku/masu/celery/tasks.py b/koku/masu/celery/tasks.py index ff22bae65c..53bc409478 100644 --- a/koku/masu/celery/tasks.py +++ b/koku/masu/celery/tasks.py @@ -319,10 +319,10 @@ def scrape_azure_storage_capacities(): disk_size_mapping = disk_fetcher.scrape_disk_size() if disk_size_mapping: for sku_prefix, disk_size in disk_size_mapping.items(): - DiskCapacity.objects.get_or_create( + DiskCapacity.objects.update_or_create( product_substring=sku_prefix, - capacity=disk_size, provider_type=Provider.PROVIDER_AZURE, + defaults={"capacity": disk_size}, ) return disk_size_mapping diff --git a/koku/masu/test/celery/test_tasks.py b/koku/masu/test/celery/test_tasks.py index f99c686f50..e09a506d84 100644 --- a/koku/masu/test/celery/test_tasks.py +++ b/koku/masu/test/celery/test_tasks.py @@ -427,6 +427,19 @@ def test_scrape_azure_storage_capacities(self, mock_celery_app): self.assertNotEqual(beforeRows, afterRows) self.assertEqual(afterRows, 14) + @patch("masu.celery.tasks.celery_app") + def test_scrape_azure_storage_capacities_updates_stale_capacity(self, mock_celery_app): + """Test that scraping updates an existing row whose capacity has changed.""" + adsf = AzureDiskSizeScraper() + DiskCapacity.objects.all().delete() + DiskCapacity.objects.create(product_substring="P80", capacity=99999, provider_type=Provider.PROVIDER_AZURE) + with requests_mock.mock() as reqmock: + reqmock.register_uri("GET", adsf.url, status_code=200, text=test_azure_scrape_output) + tasks.scrape_azure_storage_capacities() + p80 = DiskCapacity.objects.get(product_substring="P80") + self.assertEqual(p80.capacity, 32768) + self.assertEqual(DiskCapacity.objects.count(), 14) + def test_error_scrape_azure_storage_capacities(self): """Test HTTP error capture.""" adsf = AzureDiskSizeScraper()