pymongo: test_gridfs.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 gridfs package.
00018 """
00019 
00020 try:
00021     from cStringIO import StringIO
00022 except ImportError:
00023     from StringIO import StringIO
00024 import datetime
00025 import unittest
00026 import threading
00027 import time
00028 import sys
00029 sys.path[0:0] = [""]
00030 
00031 import gridfs
00032 from gridfs.errors import (FileExists,
00033                            NoFile)
00034 from test_connection import get_connection
00035 
00036 
00037 class JustWrite(threading.Thread):
00038 
00039     def __init__(self, fs):
00040         threading.Thread.__init__(self)
00041         self.fs = fs
00042 
00043     def run(self):
00044         for _ in range(10):
00045             file = self.fs.new_file(filename="test")
00046             file.write("hello")
00047             file.close()
00048 
00049 
00050 class JustRead(threading.Thread):
00051 
00052     def __init__(self, fs):
00053         threading.Thread.__init__(self)
00054         self.fs = fs
00055 
00056     def run(self):
00057         for _ in range(10):
00058             file = self.fs.get("test")
00059             assert file.read() == "hello"
00060 
00061 
00062 class TestGridfs(unittest.TestCase):
00063 
00064     def setUp(self):
00065         self.db = get_connection().pymongo_test
00066         self.db.drop_collection("fs.files")
00067         self.db.drop_collection("fs.chunks")
00068         self.db.drop_collection("alt.files")
00069         self.db.drop_collection("alt.chunks")
00070         self.fs = gridfs.GridFS(self.db)
00071         self.alt = gridfs.GridFS(self.db, "alt")
00072 
00073     def test_gridfs(self):
00074         self.assertRaises(TypeError, gridfs.GridFS, "foo")
00075         self.assertRaises(TypeError, gridfs.GridFS, self.db, 5)
00076 
00077     def test_basic(self):
00078         oid = self.fs.put("hello world")
00079         self.assertEqual("hello world", self.fs.get(oid).read())
00080         self.assertEqual(1, self.db.fs.files.count())
00081         self.assertEqual(1, self.db.fs.chunks.count())
00082 
00083         self.fs.delete(oid)
00084         self.assertRaises(NoFile, self.fs.get, oid)
00085         self.assertEqual(0, self.db.fs.files.count())
00086         self.assertEqual(0, self.db.fs.chunks.count())
00087 
00088         self.assertRaises(NoFile, self.fs.get, "foo")
00089         oid = self.fs.put("hello world", _id="foo")
00090         self.assertEqual("foo", oid)
00091         self.assertEqual("hello world", self.fs.get("foo").read())
00092 
00093 
00094     def test_list(self):
00095         self.assertEqual([], self.fs.list())
00096         self.fs.put("hello world")
00097         self.assertEqual([], self.fs.list())
00098 
00099         self.fs.put("", filename="mike")
00100         self.fs.put("foo", filename="test")
00101         self.fs.put("", filename="hello world")
00102 
00103         self.assertEqual(set(["mike", "test", "hello world"]),
00104                          set(self.fs.list()))
00105 
00106     def test_empty_file(self):
00107         oid = self.fs.put("")
00108         self.assertEqual("", self.fs.get(oid).read())
00109         self.assertEqual(1, self.db.fs.files.count())
00110         self.assertEqual(0, self.db.fs.chunks.count())
00111 
00112         raw = self.db.fs.files.find_one()
00113         self.assertEqual(0, raw["length"])
00114         self.assertEqual(oid, raw["_id"])
00115         self.assert_(isinstance(raw["uploadDate"], datetime.datetime))
00116         self.assertEqual(256*1024, raw["chunkSize"])
00117         self.assert_(isinstance(raw["md5"], basestring))
00118 
00119     def test_alt_collection(self):
00120         oid = self.alt.put("hello world")
00121         self.assertEqual("hello world", self.alt.get(oid).read())
00122         self.assertEqual(1, self.db.alt.files.count())
00123         self.assertEqual(1, self.db.alt.chunks.count())
00124 
00125         self.alt.delete(oid)
00126         self.assertRaises(NoFile, self.alt.get, oid)
00127         self.assertEqual(0, self.db.alt.files.count())
00128         self.assertEqual(0, self.db.alt.chunks.count())
00129 
00130         self.assertRaises(NoFile, self.alt.get, "foo")
00131         oid = self.alt.put("hello world", _id="foo")
00132         self.assertEqual("foo", oid)
00133         self.assertEqual("hello world", self.alt.get("foo").read())
00134 
00135         self.alt.put("", filename="mike")
00136         self.alt.put("foo", filename="test")
00137         self.alt.put("", filename="hello world")
00138 
00139         self.assertEqual(set(["mike", "test", "hello world"]),
00140                          set(self.alt.list()))
00141 
00142     def test_threaded_reads(self):
00143         self.fs.put("hello", _id="test")
00144 
00145         threads = []
00146         for i in range(10):
00147             threads.append(JustRead(self.fs))
00148             threads[i].start()
00149 
00150         for i in range(10):
00151             threads[i].join()
00152 
00153     def test_threaded_writes(self):
00154         threads = []
00155         for i in range(10):
00156             threads.append(JustWrite(self.fs))
00157             threads[i].start()
00158 
00159         for i in range(10):
00160             threads[i].join()
00161 
00162         f = self.fs.get_last_version("test")
00163         self.assertEqual(f.read(), "hello")
00164 
00165     def test_get_last_version(self):
00166         a = self.fs.put("foo", filename="test")
00167         time.sleep(0.01)
00168         b = self.fs.new_file(filename="test")
00169         b.write("bar")
00170         b.close()
00171         time.sleep(0.01)
00172         b = b._id
00173         c = self.fs.put("baz", filename="test")
00174 
00175         self.assertEqual("baz", self.fs.get_last_version("test").read())
00176         self.fs.delete(c)
00177         self.assertEqual("bar", self.fs.get_last_version("test").read())
00178         self.fs.delete(b)
00179         self.assertEqual("foo", self.fs.get_last_version("test").read())
00180         self.fs.delete(a)
00181         self.assertRaises(NoFile, self.fs.get_last_version, "test")
00182 
00183     def test_get_version(self):
00184         self.fs.put("foo", filename="test")
00185         time.sleep(0.01)
00186         self.fs.put("bar", filename="test")
00187         time.sleep(0.01)
00188         self.fs.put("baz", filename="test")
00189         time.sleep(0.01)
00190 
00191         self.assertEqual("foo", self.fs.get_version("test", 0).read())
00192         self.assertEqual("bar", self.fs.get_version("test", 1).read())
00193         self.assertEqual("baz", self.fs.get_version("test", 2).read())
00194 
00195         self.assertEqual("baz", self.fs.get_version("test", -1).read())
00196         self.assertEqual("bar", self.fs.get_version("test", -2).read())
00197         self.assertEqual("foo", self.fs.get_version("test", -3).read())
00198 
00199         self.assertRaises(NoFile, self.fs.get_version, "test", 3)
00200         self.assertRaises(NoFile, self.fs.get_version, "test", -4)
00201 
00202     def test_put_filelike(self):
00203         oid = self.fs.put(StringIO("hello world"), chunk_size=1)
00204         self.assertEqual(11, self.db.fs.chunks.count())
00205         self.assertEqual("hello world", self.fs.get(oid).read())
00206 
00207     def test_put_duplicate(self):
00208         oid = self.fs.put("hello")
00209         self.assertRaises(FileExists, self.fs.put, "world", _id=oid)
00210 
00211     def test_exists(self):
00212         oid = self.fs.put("hello")
00213         self.assert_(self.fs.exists(oid))
00214         self.assert_(self.fs.exists({"_id": oid}))
00215         self.assert_(self.fs.exists(_id=oid))
00216 
00217         self.assertFalse(self.fs.exists(filename="mike"))
00218         self.assertFalse(self.fs.exists("mike"))
00219 
00220         oid = self.fs.put("hello", filename="mike", foo=12)
00221         self.assert_(self.fs.exists(oid))
00222         self.assert_(self.fs.exists({"_id": oid}))
00223         self.assert_(self.fs.exists(_id=oid))
00224         self.assert_(self.fs.exists(filename="mike"))
00225         self.assert_(self.fs.exists({"filename": "mike"}))
00226         self.assert_(self.fs.exists(foo=12))
00227         self.assert_(self.fs.exists({"foo": 12}))
00228         self.assert_(self.fs.exists(foo={"$gt": 11}))
00229         self.assert_(self.fs.exists({"foo": {"$gt": 11}}))
00230 
00231         self.assertFalse(self.fs.exists(foo=13))
00232         self.assertFalse(self.fs.exists({"foo": 13}))
00233         self.assertFalse(self.fs.exists(foo={"$gt": 12}))
00234         self.assertFalse(self.fs.exists({"foo": {"$gt": 12}}))
00235 
00236     def test_put_unicode(self):
00237         self.assertRaises(TypeError, self.fs.put, u"hello")
00238 
00239         oid = self.fs.put(u"hello", encoding="utf-8")
00240         self.assertEqual("hello", self.fs.get(oid).read())
00241         self.assertEqual("utf-8", self.fs.get(oid).encoding)
00242 
00243         oid = self.fs.put(u"aé", encoding="iso-8859-1")
00244         self.assertEqual(u"aé".encode("iso-8859-1"), self.fs.get(oid).read())
00245         self.assertEqual("iso-8859-1", self.fs.get(oid).encoding)
00246 
00247 
00248 if __name__ == "__main__":
00249     unittest.main()