Merge pull request #232 from open-contracting/fetch_covid_active_cases

Improve fetch_covid_active_cases

.github/workflows/ci.yml

@@ -46,7 +46,7 @@ jobs:
         ./manage.py makemigrations --check --dry-run
         ./manage.py check --fail-level WARNING
         ./manage.py loaddata country redflag equitycategory equitykeywords
-        coverage run --source content,country,covidadmin,vizualization,country/management/commands manage.py test
+        coverage run --source content,country,covidadmin,vizualization,country/management/commands manage.py test --keepdb
       env:
         DB_NAME: postgres
         DB_USER: postgres

README.md

@@ -34,12 +34,32 @@ Load test data:
 ./manage.py loaddata country redflag equitycategory equitykeywords
 ```
 
+Retrieve remote data:
+
+```shell
+env DB_NAME=covid19_test ./manage.py fetch_covid_active_cases
+```
+
 ## Tasks
 
+Retrieve remote data:
+
+```shell
+./manage.py fetch_covid_active_cases
+```
+
+## Develop
+
 Run a development server:
 
 ```shell
 env DEBUG=True ./manage.py runserver
 ```
 
 To add, remove or upgrade a requirement, [follow these instructions](https://ocp-software-handbook.readthedocs.io/en/latest/python/applications.html#requirements).
+
+## Test
+
+```shell
+./manage.py test --keepdb
+```

country/management/commands/fetch_covid_active_cases.py

@@ -2,66 +2,63 @@
 import datetime
 
 import requests
+from django.contrib.postgres.aggregates import ArrayAgg
 from django.core.management.base import BaseCommand
 
 from country.models import Country, CovidMonthlyActiveCases
 
 
 class Command(BaseCommand):
-    help = "Import historical covid active cases data from covid-api.com"
+    help = "Import COVID statistics from covid-api.com"
 
     def handle(self, *args, **kwargs):
-        countries = Country.objects.all().exclude(name="Global")
-        YEARS = [2020, 2021]
+        today = datetime.date.today()
 
+        dates_in_period = set()
+        for year in range(2020, today.year + 1):
+            for month in range(1, 13):
+                day = calendar.monthrange(year, month)[1]
+                date = datetime.date(year, month, day)
+                if date > today:
+                    break
+                dates_in_period.add(date)
+
+        dates_with_data = dict(
+            Country.objects.exclude(country_code_alpha_2="gl")
+            .prefetch_related("covid_monthly_active_cases")
+            .annotate(date=ArrayAgg("covid_monthly_active_cases__covid_data_date"))
+            .values_list("country_code", "date")
+        )
+
+        countries = Country.objects.all().exclude(country_code_alpha_2="gl").order_by("country_code")
         for country in countries:
             country_code = country.country_code
-            today_date = datetime.date.today().__str__()
-            for year in YEARS:
-                for month in range(1, 12 + 1):
-                    month_end_day = calendar.monthrange(year, month)[1]
-                    date = f"{year}-{month:02}-{month_end_day:02}"
-                    date = date.split("-")
-                    date = datetime.datetime(int(date[0]), int(date[1]), int(date[2]))
-                    if today_date > str(date):
-                        date = f"{year}-{month:02}-{month_end_day:02}"
-                        existing_db_entry = CovidMonthlyActiveCases.objects.filter(
-                            country=country, covid_data_date=date
-                        ).first()
-                        if (
-                            not existing_db_entry
-                            or not existing_db_entry.active_cases_count
-                            or not existing_db_entry.death_count
-                        ):
-                            request_string = f"https://covid-api.com/api/reports?iso={country_code}&date={date}"
-                            r = requests.get(request_string)
-                            if r.status_code in [200]:
-                                covid_data = r.json()["data"]
-                                if covid_data:
-                                    print(f"Fetching {request_string}... OK")
-                                    active_cases_count = sum([province["active"] for province in covid_data])
-                                    death_count = sum([province["deaths"] for province in covid_data])
-                                    rec = CovidMonthlyActiveCases(
-                                        country=country,
-                                        covid_data_date=date,
-                                        active_cases_count=active_cases_count,
-                                        death_count=death_count,
-                                    )
-                                    rec.save()
-                                else:
-                                    print(f"Fetching {request_string}... FAILED. NO DATA")
-                                    print(r.content)
-                                    rec = CovidMonthlyActiveCases(
-                                        country=country,
-                                        covid_data_date=date,
-                                        active_cases_count=None,
-                                        death_count=None,
-                                    )
-                                    rec.save()
-                                    continue
-                            else:
-                                print(f"Fetching {request_string}... REQUEST FAILED")
-                                print(r.status_code, r.content)
-                                continue
-                        else:
-                            print(f"Data for {country_code} {date} already exists in database.")
+            dates_to_query = dates_in_period - set(dates_with_data.get(country_code, []))
+
+            for date in sorted(dates_to_query):
+                url = f"https://covid-api.com/api/reports?iso={country_code}&date={date}"
+
+                response = requests.get(url)
+                if not response.ok:
+                    self.stderr.write(f"Fetching {url}... {response.status_code}")
+                    continue
+
+                data = response.json()["data"]
+                if not data:
+                    self.stdout.write(f"Fetching {url}... NO DATA")
+                    continue
+
+                active_cases_count = 0
+                death_count = 0
+                for item in data:
+                    active_cases_count += item["active"]
+                    death_count += item["deaths"]
+
+                CovidMonthlyActiveCases(
+                    country=country,
+                    covid_data_date=date,
+                    active_cases_count=active_cases_count,
+                    death_count=death_count,
+                ).save()
+
+                self.stdout.write(f"Fetching {url}... OK")

country/tests/commands/test_fetch_covid_active_cases.py

@@ -0,0 +1,60 @@
+import datetime
+from io import StringIO
+from textwrap import dedent
+from unittest.mock import patch
+
+from django.core.management import call_command
+from django.test import TransactionTestCase
+
+from country.models import Country, CovidMonthlyActiveCases
+
+
+class CommandTests(TransactionTestCase):
+    def assertCommand(self, *args, expected_out="", expected_err="", **kwargs):
+        out = StringIO()
+        err = StringIO()
+
+        call_command(*args, **kwargs, stdout=out, stderr=out)
+
+        self.assertEqual(err.getvalue(), dedent(expected_err))
+        self.assertEqual(out.getvalue(), dedent(expected_out))
+
+    def test_command(self):
+        Country.objects.create(name="Global", country_code="GLO", country_code_alpha_2="gl", currency="USD")
+        Country.objects.create(name="Kenya", country_code="KEN", country_code_alpha_2="KE", currency="KES")
+
+        with patch("country.management.commands.fetch_covid_active_cases.datetime") as mock:
+            mock.date.today.return_value = datetime.date(2020, 4, 30)
+            mock.date.side_effect = lambda *args, **kw: datetime.date(*args, **kw)
+
+            self.assertCommand(
+                "fetch_covid_active_cases",
+                expected_out="""\
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-01-31... NO DATA
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-02-29... NO DATA
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-03-31... OK
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-04-30... OK
+                """,
+            )
+
+            self.assertEqual(
+                list(
+                    CovidMonthlyActiveCases.objects.values_list(
+                        "country__country_code", "covid_data_date", "active_cases_count", "death_count"
+                    )
+                ),
+                [
+                    ("KEN", datetime.date(2020, 3, 31), 57, 1),
+                    ("KEN", datetime.date(2020, 4, 30), 235, 17),
+                ],
+            )
+
+            self.assertCommand(
+                "fetch_covid_active_cases",
+                expected_out="""\
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-01-31... NO DATA
+                    Fetching https://covid-api.com/api/reports?iso=KEN&date=2020-02-29... NO DATA
+                """,
+            )
+
+            self.assertEqual(CovidMonthlyActiveCases.objects.count(), 2)

covidadmin/settings.py

@@ -34,7 +34,6 @@
     CELERY_BROKER_URL=(str, "pyamqp://guest@localhost/"),
     CELERY_TIMEZONE=(str, "UTC"),
     FETCH_COVID_DATA_INTERVAL=(int, 10800),
-    GOOGLE_SHEET_CREDENTIALS_JSON=(str, ""),
     FIXER_IO_API_KEY=(str, ""),
     MEDIA_URL=(str, "/media/"),
 )
@@ -83,6 +82,8 @@
     "wagtail.contrib.forms",
     "wagtail.contrib.redirects",
     "wagtail.contrib.modeladmin",
+    # See https://github.com/wagtail/wagtail/issues/1824
+    "wagtail.contrib.search_promotions",
     "wagtail.embeds",
     "wagtail.sites",
     "wagtail.users",
@@ -97,7 +98,6 @@
     "wagtail.api.v2",
     "ckeditor",
     # 'debug_toolbar',
-    # 'corsheaders',
 ]
 
 MIDDLEWARE = [
@@ -217,7 +217,6 @@
     },
 }
 
-GOOGLE_SHEET_CREDENTIALS_JSON = env("GOOGLE_SHEET_CREDENTIALS_JSON")
 FIXER_IO_API_KEY = env("FIXER_IO_API_KEY")
 
 WAGTAIL_SITE_NAME = "Covid 19 procurement explorer"