pymongo: test_grid_file.py Source File

00001 # -*- coding: utf-8 -*-
00002 #
00003 # Copyright 2009-2010 10gen, Inc.
00004 #
00005 # Licensed under the Apache License, Version 2.0 (the "License");
00006 # you may not use this file except in compliance with the License.
00007 # You may obtain a copy of the License at
00008 #
00009 # http://www.apache.org/licenses/LICENSE-2.0
00010 #
00011 # Unless required by applicable law or agreed to in writing, software
00012 # distributed under the License is distributed on an "AS IS" BASIS,
00013 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00014 # See the License for the specific language governing permissions and
00015 # limitations under the License.
00016 
00017 """Tests for the grid_file module.
00018 """
00019 
00020 try:
00021     from cStringIO import StringIO
00022 except ImportError:
00023     from StringIO import StringIO
00024 import datetime
00025 import os
00026 import sys
00027 import unittest
00028 sys.path[0:0] = [""]
00029 
00030 from bson.objectid import ObjectId
00031 from gridfs.grid_file import (_SEEK_CUR,
00032                               _SEEK_END,
00033                               GridIn,
00034                               GridFile,
00035                               GridOut)
00036 from gridfs.errors import (NoFile,
00037                            UnsupportedAPI)
00038 from test_connection import get_connection
00039 import qcheck
00040 
00041 
00042 class TestGridFile(unittest.TestCase):
00043 
00044     def setUp(self):
00045         self.db = get_connection().pymongo_test
00046         self.db.fs.files.remove({})
00047         self.db.fs.chunks.remove({})
00048 
00049     def test_basic(self):
00050         f = GridIn(self.db.fs, filename="test")
00051         f.write("hello world")
00052         f.close()
00053         self.assertEqual(1, self.db.fs.files.find().count())
00054         self.assertEqual(1, self.db.fs.chunks.find().count())
00055 
00056         g = GridOut(self.db.fs, f._id)
00057         self.assertEqual("hello world", g.read())
00058 
00059         # make sure it's still there...
00060         g = GridOut(self.db.fs, f._id)
00061         self.assertEqual("hello world", g.read())
00062 
00063 
00064         f = GridIn(self.db.fs, filename="test")
00065         f.close()
00066         self.assertEqual(2, self.db.fs.files.find().count())
00067         self.assertEqual(1, self.db.fs.chunks.find().count())
00068 
00069         g = GridOut(self.db.fs, f._id)
00070         self.assertEqual("", g.read())
00071 
00072     def test_md5(self):
00073         f = GridIn(self.db.fs)
00074         f.write("hello world\n")
00075         f.close()
00076         self.assertEqual("6f5902ac237024bdd0c176cb93063dc4", f.md5)
00077 
00078     def test_alternate_collection(self):
00079         self.db.alt.files.remove({})
00080         self.db.alt.chunks.remove({})
00081 
00082         f = GridIn(self.db.alt)
00083         f.write("hello world")
00084         f.close()
00085 
00086         self.assertEqual(1, self.db.alt.files.find().count())
00087         self.assertEqual(1, self.db.alt.chunks.find().count())
00088 
00089         g = GridOut(self.db.alt, f._id)
00090         self.assertEqual("hello world", g.read())
00091 
00092         # test that md5 still works...
00093         self.assertEqual("5eb63bbbe01eeed093cb22bb8f5acdc3", g.md5)
00094 
00095     def test_grid_file(self):
00096         self.assertRaises(UnsupportedAPI, GridFile)
00097 
00098     def test_grid_in_default_opts(self):
00099         self.assertRaises(TypeError, GridIn, "foo")
00100 
00101         a = GridIn(self.db.fs)
00102 
00103         self.assert_(isinstance(a._id, ObjectId))
00104         self.assertRaises(AttributeError, setattr, a, "_id", 5)
00105 
00106         self.assertEqual(None, a.filename)
00107         a.filename = "my_file"
00108         self.assertEqual("my_file", a.filename)
00109 
00110         self.assertEqual(None, a.content_type)
00111         a.content_type = "text/html"
00112         self.assertEqual("text/html", a.content_type)
00113 
00114         self.assertRaises(AttributeError, getattr, a, "length")
00115         self.assertRaises(AttributeError, setattr, a, "length", 5)
00116 
00117         self.assertEqual(256 * 1024, a.chunk_size)
00118         self.assertRaises(AttributeError, setattr, a, "chunk_size", 5)
00119 
00120         self.assertRaises(AttributeError, getattr, a, "upload_date")
00121         self.assertRaises(AttributeError, setattr, a, "upload_date", 5)
00122 
00123         self.assertRaises(AttributeError, getattr, a, "aliases")
00124         a.aliases = ["foo"]
00125         self.assertEqual(["foo"], a.aliases)
00126 
00127         self.assertRaises(AttributeError, getattr, a, "metadata")
00128         a.metadata = {"foo": 1}
00129         self.assertEqual({"foo": 1}, a.metadata)
00130 
00131         self.assertRaises(AttributeError, getattr, a, "md5")
00132         self.assertRaises(AttributeError, setattr, a, "md5", 5)
00133 
00134         a.close()
00135 
00136         self.assert_(isinstance(a._id, ObjectId))
00137         self.assertRaises(AttributeError, setattr, a, "_id", 5)
00138 
00139         self.assertEqual("my_file", a.filename)
00140 
00141         self.assertEqual("text/html", a.content_type)
00142 
00143         self.assertEqual(0, a.length)
00144         self.assertRaises(AttributeError, setattr, a, "length", 5)
00145 
00146         self.assertEqual(256 * 1024, a.chunk_size)
00147         self.assertRaises(AttributeError, setattr, a, "chunk_size", 5)
00148 
00149         self.assert_(isinstance(a.upload_date, datetime.datetime))
00150         self.assertRaises(AttributeError, setattr, a, "upload_date", 5)
00151 
00152         self.assertEqual(["foo"], a.aliases)
00153 
00154         self.assertEqual({"foo": 1}, a.metadata)
00155 
00156         self.assertEqual("d41d8cd98f00b204e9800998ecf8427e", a.md5)
00157         self.assertRaises(AttributeError, setattr, a, "md5", 5)
00158 
00159     def test_grid_in_custom_opts(self):
00160         self.assertRaises(TypeError, GridIn, "foo")
00161 
00162         a = GridIn(self.db.fs, _id=5, filename="my_file",
00163                    contentType="text/html", chunkSize=1000, aliases=["foo"],
00164                    metadata={"foo": 1, "bar": 2}, bar=3, baz="hello")
00165 
00166         self.assertEqual(5, a._id)
00167         self.assertEqual("my_file", a.filename)
00168         self.assertEqual("text/html", a.content_type)
00169         self.assertEqual(1000, a.chunk_size)
00170         self.assertEqual(["foo"], a.aliases)
00171         self.assertEqual({"foo": 1, "bar": 2}, a.metadata)
00172         self.assertEqual(3, a.bar)
00173         self.assertEqual("hello", a.baz)
00174         self.assertRaises(AttributeError, getattr, a, "mike")
00175 
00176         b = GridIn(self.db.fs, content_type="text/html", chunk_size=1000, baz=100)
00177         self.assertEqual("text/html", b.content_type)
00178         self.assertEqual(1000, b.chunk_size)
00179         self.assertEqual(100, b.baz)
00180 
00181     def test_grid_out_default_opts(self):
00182         self.assertRaises(TypeError, GridOut, "foo")
00183 
00184         self.assertRaises(NoFile, GridOut, self.db.fs, 5)
00185 
00186         a = GridIn(self.db.fs)
00187         a.close()
00188 
00189         b = GridOut(self.db.fs, a._id)
00190 
00191         self.assertEqual(a._id, b._id)
00192         self.assertEqual(0, b.length)
00193         self.assertEqual(None, b.content_type)
00194         self.assertEqual(256 * 1024, b.chunk_size)
00195         self.assert_(isinstance(b.upload_date, datetime.datetime))
00196         self.assertEqual(None, b.aliases)
00197         self.assertEqual(None, b.metadata)
00198         self.assertEqual("d41d8cd98f00b204e9800998ecf8427e", b.md5)
00199 
00200         for attr in ["_id", "name", "content_type", "length", "chunk_size",
00201                      "upload_date", "aliases", "metadata", "md5"]:
00202             self.assertRaises(AttributeError, setattr, b, attr, 5)
00203 
00204     def test_grid_out_custom_opts(self):
00205         a = GridIn(self.db.fs, _id=5, filename="my_file",
00206                    contentType="text/html", chunkSize=1000, aliases=["foo"],
00207                    metadata={"foo": 1, "bar": 2}, bar=3, baz="hello")
00208         a.write("hello world")
00209         a.close()
00210 
00211         b = GridOut(self.db.fs, 5)
00212 
00213         self.assertEqual(5, b._id)
00214         self.assertEqual(11, b.length)
00215         self.assertEqual("text/html", b.content_type)
00216         self.assertEqual(1000, b.chunk_size)
00217         self.assert_(isinstance(b.upload_date, datetime.datetime))
00218         self.assertEqual(["foo"], b.aliases)
00219         self.assertEqual({"foo": 1, "bar": 2}, b.metadata)
00220         self.assertEqual(3, b.bar)
00221         self.assertEqual("5eb63bbbe01eeed093cb22bb8f5acdc3", b.md5)
00222 
00223         for attr in ["_id", "name", "content_type", "length", "chunk_size",
00224                      "upload_date", "aliases", "metadata", "md5"]:
00225             self.assertRaises(AttributeError, setattr, b, attr, 5)
00226 
00227     def test_grid_out_file_document(self):
00228         a = GridIn(self.db.fs)
00229         a.write("foo bar")
00230         a.close()
00231 
00232         b = GridOut(self.db.fs, file_document=self.db.fs.files.find_one())
00233         self.assertEqual("foo bar", b.read())
00234 
00235         c = GridOut(self.db.fs, 5, file_document=self.db.fs.files.find_one())
00236         self.assertEqual("foo bar", c.read())
00237 
00238         self.assertRaises(NoFile, GridOut, self.db.fs, file_document={})
00239 
00240     def test_write_file_like(self):
00241         a = GridIn(self.db.fs)
00242         a.write("hello world")
00243         a.close()
00244 
00245         b = GridOut(self.db.fs, a._id)
00246 
00247         c = GridIn(self.db.fs)
00248         c.write(b)
00249         c.close()
00250 
00251         d = GridOut(self.db.fs, c._id)
00252         self.assertEqual("hello world", d.read())
00253 
00254         e = GridIn(self.db.fs, chunk_size=2)
00255         e.write("hello")
00256         buffer = StringIO(" world")
00257         e.write(buffer)
00258         e.write(" and mongodb")
00259         e.close()
00260         self.assertEqual("hello world and mongodb",
00261                          GridOut(self.db.fs, e._id).read())
00262 
00263     def test_write_lines(self):
00264         a = GridIn(self.db.fs)
00265         a.writelines(["hello ", "world"])
00266         a.close()
00267 
00268         self.assertEqual("hello world", GridOut(self.db.fs, a._id).read())
00269 
00270     def test_close(self):
00271         f = GridIn(self.db.fs)
00272         f.close()
00273         self.assertRaises(ValueError, f.write, "test")
00274         f.close()
00275 
00276     def test_multi_chunk_file(self):
00277         random_string = qcheck.gen_string(qcheck.lift(300000))()
00278 
00279         f = GridIn(self.db.fs)
00280         f.write(random_string)
00281         f.close()
00282 
00283         self.assertEqual(1, self.db.fs.files.find().count())
00284         self.assertEqual(2, self.db.fs.chunks.find().count())
00285 
00286         g = GridOut(self.db.fs, f._id)
00287         self.assertEqual(random_string, g.read())
00288 
00289     def test_small_chunks(self):
00290         self.files = 0
00291         self.chunks = 0
00292 
00293         def helper(data):
00294             f = GridIn(self.db.fs, chunkSize=1)
00295             f.write(data)
00296             f.close()
00297 
00298             self.files += 1
00299             self.chunks += len(data)
00300 
00301             self.assertEqual(self.files, self.db.fs.files.find().count())
00302             self.assertEqual(self.chunks, self.db.fs.chunks.find().count())
00303 
00304             g = GridOut(self.db.fs, f._id)
00305             self.assertEqual(data, g.read())
00306 
00307             g = GridOut(self.db.fs, f._id)
00308             self.assertEqual(data, g.read(10) + g.read(10))
00309             return True
00310 
00311         qcheck.check_unittest(self, helper,
00312                               qcheck.gen_string(qcheck.gen_range(0, 20)))
00313 
00314     def test_seek(self):
00315         f = GridIn(self.db.fs, chunkSize=3)
00316         f.write("hello world")
00317         f.close()
00318 
00319         g = GridOut(self.db.fs, f._id)
00320         self.assertEqual("hello world", g.read())
00321         g.seek(0)
00322         self.assertEqual("hello world", g.read())
00323         g.seek(1)
00324         self.assertEqual("ello world", g.read())
00325         self.assertRaises(IOError, g.seek, -1)
00326 
00327         g.seek(-3, _SEEK_END)
00328         self.assertEqual("rld", g.read())
00329         g.seek(0, _SEEK_END)
00330         self.assertEqual("", g.read())
00331         self.assertRaises(IOError, g.seek, -100, _SEEK_END)
00332 
00333         g.seek(3)
00334         g.seek(3, _SEEK_CUR)
00335         self.assertEqual("world", g.read())
00336         self.assertRaises(IOError, g.seek, -100, _SEEK_CUR)
00337 
00338     def test_tell(self):
00339         f = GridIn(self.db.fs, chunkSize=3)
00340         f.write("hello world")
00341         f.close()
00342 
00343         g = GridOut(self.db.fs, f._id)
00344         self.assertEqual(0, g.tell())
00345         g.read(0)
00346         self.assertEqual(0, g.tell())
00347         g.read(1)
00348         self.assertEqual(1, g.tell())
00349         g.read(2)
00350         self.assertEqual(3, g.tell())
00351         g.read()
00352         self.assertEqual(g.length, g.tell())
00353 
00354     def test_multiple_reads(self):
00355         f = GridIn(self.db.fs, chunkSize=3)
00356         f.write("hello world")
00357         f.close()
00358 
00359         g = GridOut(self.db.fs, f._id)
00360         self.assertEqual("he", g.read(2))
00361         self.assertEqual("ll", g.read(2))
00362         self.assertEqual("o ", g.read(2))
00363         self.assertEqual("wo", g.read(2))
00364         self.assertEqual("rl", g.read(2))
00365         self.assertEqual("d", g.read(2))
00366         self.assertEqual("", g.read(2))
00367 
00368     def test_readline(self):
00369         f = GridIn(self.db.fs, chunkSize=5)
00370         f.write("""Hello world,
00371 How are you?
00372 Hope all is well.
00373 Bye""")
00374         f.close()
00375 
00376         g = GridOut(self.db.fs, f._id)
00377         self.assertEqual("H", g.read(1))
00378         self.assertEqual("ello world,\n", g.readline())
00379         self.assertEqual("How a", g.readline(5))
00380         self.assertEqual("", g.readline(0))
00381         self.assertEqual("re you?\n", g.readline())
00382         self.assertEqual("Hope all is well.\n", g.readline(1000))
00383         self.assertEqual("Bye", g.readline())
00384         self.assertEqual("", g.readline())
00385 
00386     def test_iterator(self):
00387         f = GridIn(self.db.fs)
00388         f.close()
00389         g = GridOut(self.db.fs, f._id)
00390         self.assertEqual([], list(g))
00391 
00392         f = GridIn(self.db.fs)
00393         f.write("hello world")
00394         f.close()
00395         g = GridOut(self.db.fs, f._id)
00396         self.assertEqual(["hello world"], list(g))
00397         self.assertEqual("hello", g.read(5))
00398         self.assertEqual(["hello world"], list(g))
00399         self.assertEqual(" worl", g.read(5))
00400 
00401         f = GridIn(self.db.fs, chunk_size=2)
00402         f.write("hello world")
00403         f.close()
00404         g = GridOut(self.db.fs, f._id)
00405         self.assertEqual(["he", "ll", "o ", "wo", "rl", "d"], list(g))
00406 
00407     def test_read_chunks_unaligned_buffer_size(self):
00408         in_data = "This is a text that doesn't quite fit in a single 16-byte chunk."
00409         f = GridIn(self.db.fs, chunkSize=16)
00410         f.write(in_data)
00411         f.close()
00412 
00413         g = GridOut(self.db.fs, f._id)
00414         out_data = ''
00415         while 1:
00416             s = g.read(13)
00417             if not s:
00418                 break
00419             out_data += s
00420 
00421         self.assertEqual(in_data, out_data)
00422 
00423     def test_write_unicode(self):
00424         f = GridIn(self.db.fs)
00425         self.assertRaises(TypeError, f.write, u"foo")
00426 
00427         f = GridIn(self.db.fs, encoding="utf-8")
00428         f.write(u"foo")
00429         f.close()
00430 
00431         g = GridOut(self.db.fs, f._id)
00432         self.assertEqual("foo", g.read())
00433 
00434         f = GridIn(self.db.fs, encoding="iso-8859-1")
00435         f.write(u"aé")
00436         f.close()
00437 
00438         g = GridOut(self.db.fs, f._id)
00439         self.assertEqual(u"aé".encode("iso-8859-1"), g.read())
00440 
00441     def test_set_after_close(self):
00442         f = GridIn(self.db.fs, _id="foo", bar="baz")
00443 
00444         self.assertEqual("foo", f._id)
00445         self.assertEqual("baz", f.bar)
00446         self.assertRaises(AttributeError, getattr, f, "baz")
00447         self.assertRaises(AttributeError, getattr, f, "uploadDate")
00448 
00449         self.assertRaises(AttributeError, setattr, f, "_id", 5)
00450         f.bar = "foo"
00451         f.baz = 5
00452 
00453         self.assertEqual("foo", f._id)
00454         self.assertEqual("foo", f.bar)
00455         self.assertEqual(5, f.baz)
00456         self.assertRaises(AttributeError, getattr, f, "uploadDate")
00457 
00458         f.close()
00459 
00460         self.assertEqual("foo", f._id)
00461         self.assertEqual("foo", f.bar)
00462         self.assertEqual(5, f.baz)
00463         self.assert_(f.uploadDate)
00464 
00465         self.assertRaises(AttributeError, setattr, f, "_id", 5)
00466         f.bar = "a"
00467         f.baz = "b"
00468         self.assertRaises(AttributeError, setattr, f, "upload_date", 5)
00469 
00470         g = GridOut(self.db.fs, f._id)
00471         self.assertEqual("a", f.bar)
00472         self.assertEqual("b", f.baz)
00473 
00474 
00475 if __name__ == "__main__":
00476     unittest.main()