test GridFS helpers (and fix some bugs along the way)

Author: dcrosta

flask_pymongo/__init__.py

@@ -211,16 +211,16 @@ def get_upload(filename):
 
         # mostly copied from flask/helpers.py, with
         # modifications for GridFS
-        data = wrap_file(request.environ, fileobj, buffer_size=1024 * 256)
+        data = wrap_file(request.environ, fileobj, buffer_size=1024 * 255)
         response = current_app.response_class(
             data,
             mimetype=fileobj.content_type,
-            direct_passthrough=True)
+            direct_passthrough=True,
+        )
         response.content_length = fileobj.length
         response.last_modified = fileobj.upload_date
         response.set_etag(fileobj.md5)
         response.cache_control.max_age = cache_for
-        response.cache_control.s_max_age = cache_for
         response.cache_control.public = True
         response.make_conditional(request)
         return response

flask_pymongo/tests/test_gridfs.py

@@ -0,0 +1,83 @@
+from hashlib import md5
+from io import BytesIO
+
+from gridfs import GridFS
+from werkzeug.exceptions import NotFound
+import pytest
+
+from flask_pymongo.tests.util import FlaskPyMongoTest
+
+
+class GridFSCleanupMixin(object):
+
+    def tearDown(self):
+        gridfs = GridFS(self.mongo.db)
+        files = list(gridfs.find())
+        for gridfile in files:
+            gridfs.delete(gridfile._id)
+
+        super(GridFSCleanupMixin, self).tearDown()
+
+
+class TestSaveFile(GridFSCleanupMixin, FlaskPyMongoTest):
+
+    def test_it_saves_files(self):
+        fileobj = BytesIO(b"these are the bytes")
+
+        self.mongo.save_file("my-file", fileobj)
+
+        gridfs = GridFS(self.mongo.db)
+        assert gridfs.exists({"filename": "my-file"})
+
+    def test_it_guesses_type_from_filename(self):
+        fileobj = BytesIO(b"these are the bytes")
+
+        self.mongo.save_file("my-file.txt", fileobj)
+
+        gridfs = GridFS(self.mongo.db)
+        gridfile = gridfs.find_one({"filename": "my-file.txt"})
+        assert gridfile.content_type == "text/plain"
+
+
+class TestSendFile(GridFSCleanupMixin, FlaskPyMongoTest):
+
+    def setUp(self):
+        super(TestSendFile, self).setUp()
+
+        # make it bigger than 1 gridfs chunk
+        self.myfile = BytesIO(b"a" * 500 * 1024)
+        self.mongo.save_file("myfile.txt", self.myfile)
+
+    def test_it_404s_for_missing_files(self):
+        with pytest.raises(NotFound):
+            self.mongo.send_file("no-such-file.txt")
+
+    def test_it_sets_content_type(self):
+        resp = self.mongo.send_file("myfile.txt")
+        assert resp.content_type.startswith("text/plain")
+
+    def test_it_sets_content_length(self):
+        resp = self.mongo.send_file("myfile.txt")
+        assert resp.content_length == len(self.myfile.getvalue())
+
+    def test_it_sets_supports_conditional_gets(self):
+        # a basic conditional GET
+        environ_args = {
+            "method": "GET",
+            "headers": {
+                "If-None-Match": md5(self.myfile.getvalue()).hexdigest(),
+            },
+        }
+
+        with self.app.test_request_context(**environ_args):
+            resp = self.mongo.send_file("myfile.txt")
+            assert resp.status_code == 304
+
+    def test_it_sets_cache_headers(self):
+        resp = self.mongo.send_file("myfile.txt", cache_for=60)
+        assert resp.cache_control.max_age == 60
+        assert resp.cache_control.public is True
+
+    def test_it_streams_results(self):
+        resp = self.mongo.send_file("myfile.txt")
+        assert resp.is_streamed