Support local file paths as a value of PEP_ARTIFACT_URL (#1428)

mozillazg · berkerpeksag · commit f17ceb834350 · 2019-04-21T18:07:13.000+03:00
diff --git a/peps/management/commands/generate_pep_pages.py b/peps/management/commands/generate_pep_pages.py
@@ -52,20 +52,7 @@ def verbose(msg):
 
         with ExitStack() as stack:
             verbose(f"== Fetching PEP artifact from {settings.PEP_ARTIFACT_URL}")
-            peps_last_updated = get_peps_last_updated()
-            with requests.get(settings.PEP_ARTIFACT_URL, stream=True) as r:
-                artifact_last_modified = parsedate(r.headers['last-modified'])
-                if peps_last_updated > artifact_last_modified:
-                    verbose(f"== No update to artifacts, we're done here!")
-                    return
-
-                temp_file = stack.enter_context(TemporaryFile())
-                for chunk in r.iter_content(chunk_size=8192):
-                    if chunk:
-                        temp_file.write(chunk)
-
-            temp_file.seek(0)
-
+            temp_file = self.get_artifact_tarball(stack, verbose)
             temp_dir = stack.enter_context(TemporaryDirectory())
             tar_ball = stack.enter_context(TarFile.open(fileobj=temp_file, mode='r:gz'))
             tar_ball.extractall(path=temp_dir, numeric_owner=False)
@@ -129,3 +116,23 @@ def verbose(msg):
                     verbose("- Skipping non-PEP related image '{}'".format(img))
 
         verbose("== Finished")
+
+    def get_artifact_tarball(self, stack, verbose):
+        artifact_url = settings.PEP_ARTIFACT_URL
+        if not artifact_url.startswith(('http://', 'https://')):
+            return stack.enter_context(open(artifact_url, 'rb'))
+
+        peps_last_updated = get_peps_last_updated()
+        with requests.get(artifact_url, stream=True) as r:
+            artifact_last_modified = parsedate(r.headers['last-modified'])
+            if peps_last_updated > artifact_last_modified:
+                verbose(f"== No update to artifacts, we're done here!")
+                return
+
+            temp_file = stack.enter_context(TemporaryFile())
+            for chunk in r.iter_content(chunk_size=8192):
+                if chunk:
+                    temp_file.write(chunk)
+
+        temp_file.seek(0)
+        return temp_file
diff --git a/peps/tests/test_commands.py b/peps/tests/test_commands.py
@@ -32,7 +32,11 @@ def setUp(self):
         )
 
     @responses.activate
-    def test_generate_pep_pages_real(self):
+    def test_generate_pep_pages_real_with_remote_artifact(self):
+        call_command('generate_pep_pages')
+
+    @override_settings(PEP_ARTIFACT_URL=FAKE_PEP_ARTIFACT)
+    def test_generate_pep_pages_real_with_local_artifact(self):
         call_command('generate_pep_pages')
 
     @responses.activate
diff --git a/pydotorg/settings/local.py b/pydotorg/settings/local.py
@@ -24,8 +24,9 @@
 
 EMAIL_BACKEND = 'django.core.mail.backends.console.EmailBackend'
 
-# Set the URL to where to fetch PEP artifacts from
-PEP_ARTIFACT_URL = 'https://pythondotorg-assets-staging.s3.amazonaws.com/fake-peps.tar.gz'
+# Set the path to where to fetch PEP artifacts from.
+# The value can be a local path or a remote URL.
+PEP_ARTIFACT_URL = os.path.join(BASE, 'peps/tests/peps.tar.gz')
 
 # Use Dummy SASS compiler to avoid performance issues and remove the need to
 # have a sass compiler installed at all during local development if you aren't
