Added command to import all starships

prdickson · prdickson · commit 2150c213dfde · 2018-05-19T13:16:35.000+01:00
diff --git a/requirements.in b/requirements.in
@@ -4,3 +4,4 @@ psycopg2>=2.6.2,<2.7
 djangorestframework
 django-filter
 django-extensions
+requests
diff --git a/shiptrader/management/__init__.py b/shiptrader/management/__init__.py
diff --git a/shiptrader/management/commands/__init__.py b/shiptrader/management/commands/__init__.py
diff --git a/shiptrader/management/commands/import_starships.py b/shiptrader/management/commands/import_starships.py
@@ -0,0 +1,37 @@
+import requests
+from django.core.management.base import BaseCommand, CommandError
+
+from shiptrader.models import Starship
+
+
+class Command(BaseCommand):
+    help = 'Retrieves starship data from the swapi.co api '
+
+    def handle(self, *args, **options):
+
+        url = 'https://swapi.co/api/starships/'
+        starships = []
+
+        while url is not None:
+            response = requests.get(url)
+            response_json = response.json()
+            url = response_json['next']
+            starships += [build_starship(result) for result in response_json['results']]
+
+        Starship.objects.bulk_create(starships)
+
+
+def build_starship(data):
+    return Starship(name=data['name'],
+                    model=data['model'],
+                    starship_class=data['starship_class'],
+                    manufacturer=data['manufacturer'],
+                    length=float(convert_unknown(data['length'])),
+                    hyperdrive_rating=float(convert_unknown(data['hyperdrive_rating'])),
+                    cargo_capacity=int(convert_unknown(data['cargo_capacity'])),
+                    crew=int(convert_unknown(data['crew'])),
+                    passengers=int(convert_unknown(data['passengers'])))
+
+
+def convert_unknown(value):
+    return value.replace(',', '.') if value != 'unknown' else 0
diff --git a/shiptrader/tests/test_import_starships.py b/shiptrader/tests/test_import_starships.py
@@ -0,0 +1,28 @@
+from django.core.management import call_command
+from django.test import TestCase
+from rest_framework.reverse import reverse
+
+from shiptrader.models import Starship
+
+
+class TestImportStarships(TestCase):
+
+    def test_can_import_starships(self):
+
+        call_command('import_starships')
+        starships = Starship.objects.all()
+        self.assertGreater(len(starships), 0)
+
+        executor = next(filter(lambda s: s.name == 'Executor', starships))
+
+        self.assertEqual(executor.name, 'Executor')
+        self.assertEqual(executor.model, 'Executor-class star dreadnought')
+        self.assertEqual(executor.starship_class, 'Star dreadnought')
+        self.assertEqual(executor.manufacturer, 'Kuat Drive Yards, Fondor Shipyards')
+
+        self.assertEqual(executor.length, 19000.0)
+        self.assertEqual(executor.hyperdrive_rating, 2.0)
+        self.assertEqual(executor.cargo_capacity, 250000000)
+
+        self.assertEqual(executor.crew, 279144)
+        self.assertEqual(executor.passengers, 38000)
diff --git a/shiptrader/views.py b/shiptrader/views.py
@@ -10,5 +10,7 @@ class Meta:
 
 
 class StarshipViewSet(viewsets.ReadOnlyModelViewSet):
+    """View that allows all starships to be listed"""
+
     queryset = Starship.objects.all()
     serializer_class = StarshipSerializer
