summaryrefslogtreecommitdiff
path: root/tests/test_lobject.py
blob: f4a684fdd2a8c3348d37a00d137de7c74af672be (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
#!/usr/bin/env python
import os
import shutil
import tempfile
import unittest
import warnings

import psycopg2
import psycopg2.extensions
import tests


class LargeObjectTests(unittest.TestCase):

    def setUp(self):
        self.conn = psycopg2.connect(tests.dsn)
        self.lo_oid = None
        self.tmpdir = None

    def tearDown(self):
        if self.tmpdir:
            shutil.rmtree(self.tmpdir, ignore_errors=True)
        if self.lo_oid is not None:
            self.conn.rollback()
            try:
                lo = self.conn.lobject(self.lo_oid, "n")
            except psycopg2.OperationalError:
                pass
            else:
                lo.unlink()
        self.conn.close()

    def test_create(self):
        lo = self.conn.lobject()
        self.assertNotEqual(lo, None)
        self.assertEqual(lo.mode, "w")

    def test_open_non_existent(self):
        # By creating then removing a large object, we get an Oid that
        # should be unused.
        lo = self.conn.lobject()
        lo.unlink()
        self.assertRaises(psycopg2.OperationalError, self.conn.lobject, lo.oid)

    def test_open_existing(self):
        lo = self.conn.lobject()
        lo2 = self.conn.lobject(lo.oid)
        self.assertNotEqual(lo2, None)
        self.assertEqual(lo2.oid, lo.oid)
        self.assertEqual(lo2.mode, "r")

    def test_open_for_write(self):
        lo = self.conn.lobject()
        lo2 = self.conn.lobject(lo.oid, "w")
        self.assertEqual(lo2.mode, "w")
        lo2.write("some data")

    def test_open_mode_n(self):
        # Openning an object in mode "n" gives us a closed lobject.
        lo = self.conn.lobject()
        lo.close()

        lo2 = self.conn.lobject(lo.oid, "n")
        self.assertEqual(lo2.oid, lo.oid)
        self.assertEqual(lo2.closed, True)

    def test_create_with_oid(self):
        # Create and delete a large object to get an unused Oid.
        lo = self.conn.lobject()
        oid = lo.oid
        lo.unlink()

        lo = self.conn.lobject(0, "w", oid)
        self.assertEqual(lo.oid, oid)

    def test_create_with_existing_oid(self):
        lo = self.conn.lobject()
        lo.close()

        self.assertRaises(psycopg2.OperationalError,
                          self.conn.lobject, 0, "w", lo.oid)

    def test_import(self):
        self.tmpdir = tempfile.mkdtemp()
        filename = os.path.join(self.tmpdir, "data.txt")
        fp = open(filename, "wb")
        fp.write("some data")
        fp.close()

        lo = self.conn.lobject(0, "r", 0, filename)
        self.assertEqual(lo.read(), "some data")

    def test_close(self):
        lo = self.conn.lobject()
        self.assertEqual(lo.closed, False)
        lo.close()
        self.assertEqual(lo.closed, True)

    def test_write(self):
        lo = self.conn.lobject()
        self.assertEqual(lo.write("some data"), len("some data"))

    def test_write_large(self):
        lo = self.conn.lobject()
        data = "data" * 1000000
        self.assertEqual(lo.write(data), len(data))

    def test_read(self):
        lo = self.conn.lobject()
        length = lo.write("some data")
        lo.close()

        lo = self.conn.lobject(lo.oid)
        self.assertEqual(lo.read(4), "some")
        self.assertEqual(lo.read(), " data")

    def test_read_large(self):
        lo = self.conn.lobject()
        data = "data" * 1000000
        length = lo.write("some"+data)
        lo.close()

        lo = self.conn.lobject(lo.oid)
        self.assertEqual(lo.read(4), "some")
        self.assertEqual(lo.read(), data)

    def test_seek_tell(self):
        lo = self.conn.lobject()
        length = lo.write("some data")
        self.assertEqual(lo.tell(), length)
        lo.close()
        lo = self.conn.lobject(lo.oid)

        self.assertEqual(lo.seek(5, 0), 5)
        self.assertEqual(lo.tell(), 5)
        self.assertEqual(lo.read(), "data")

        # SEEK_CUR: relative current location
        lo.seek(5)
        self.assertEqual(lo.seek(2, 1), 7)
        self.assertEqual(lo.tell(), 7)
        self.assertEqual(lo.read(), "ta")

        # SEEK_END: relative to end of file
        self.assertEqual(lo.seek(-2, 2), length - 2)
        self.assertEqual(lo.read(), "ta")

    def test_unlink(self):
        lo = self.conn.lobject()
        lo.unlink()

        # the object doesn't exist now, so we can't reopen it.
        self.assertRaises(psycopg2.OperationalError, self.conn.lobject, lo.oid)
        # And the object has been closed.
        self.assertEquals(lo.closed, True)

    def test_export(self):
        lo = self.conn.lobject()
        lo.write("some data")

        self.tmpdir = tempfile.mkdtemp()
        filename = os.path.join(self.tmpdir, "data.txt")
        lo.export(filename)
        self.assertTrue(os.path.exists(filename))
        self.assertEqual(open(filename, "rb").read(), "some data")

    def test_close_twice(self):
        lo = self.conn.lobject()
        lo.close()
        lo.close()

    def test_write_after_close(self):
        lo = self.conn.lobject()
        lo.close()
        self.assertRaises(psycopg2.InterfaceError, lo.write, "some data")

    def test_read_after_close(self):
        lo = self.conn.lobject()
        lo.close()
        self.assertRaises(psycopg2.InterfaceError, lo.read, 5)

    def test_seek_after_close(self):
        lo = self.conn.lobject()
        lo.close()
        self.assertRaises(psycopg2.InterfaceError, lo.seek, 0)

    def test_tell_after_close(self):
        lo = self.conn.lobject()
        lo.close()
        self.assertRaises(psycopg2.InterfaceError, lo.tell)

    def test_unlink_after_close(self):
        lo = self.conn.lobject()
        lo.close()
        # Unlink works on closed files.
        lo.unlink()

    def test_export_after_close(self):
        lo = self.conn.lobject()
        lo.write("some data")
        lo.close()

        self.tmpdir = tempfile.mkdtemp()
        filename = os.path.join(self.tmpdir, "data.txt")
        lo.export(filename)
        self.assertTrue(os.path.exists(filename))
        self.assertEqual(open(filename, "rb").read(), "some data")

    def test_close_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        # Closing outside of the transaction is okay.
        lo.close()

    def test_write_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        self.assertRaises(psycopg2.ProgrammingError, lo.write, "some data")

    def test_read_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        self.assertRaises(psycopg2.ProgrammingError, lo.read, 5)

    def test_seek_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        self.assertRaises(psycopg2.ProgrammingError, lo.seek, 0)

    def test_tell_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        self.assertRaises(psycopg2.ProgrammingError, lo.tell)

    def test_unlink_after_commit(self):
        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        # Unlink of stale lobject is okay
        lo.unlink()

    def test_export_after_commit(self):
        lo = self.conn.lobject()
        lo.write("some data")
        self.conn.commit()

        self.tmpdir = tempfile.mkdtemp()
        filename = os.path.join(self.tmpdir, "data.txt")
        lo.export(filename)
        self.assertTrue(os.path.exists(filename))
        self.assertEqual(open(filename, "rb").read(), "some data")


class LargeObjectTruncateTests(LargeObjectTests):

    skip = None

    def setUp(self):
        LargeObjectTests.setUp(self)

        if self.skip is None:
            self.skip = False
            if self.conn.server_version < 80300:
                warnings.warn("Large object truncate tests skipped, "
                              "the server does not support them")
                self.skip = True

            if not hasattr(psycopg2.extensions.lobject, 'truncate'):
                warnings.warn("Large object truncate tests skipped, "
                              "psycopg2 has been built against an old library")
                self.skip = True

    def test_truncate(self):
        if self.skip:
            return

        lo = self.conn.lobject()
        lo.write("some data")
        lo.close()

        lo = self.conn.lobject(lo.oid, "w")
        lo.truncate(4)

        # seek position unchanged
        self.assertEqual(lo.tell(), 0)
        # data truncated
        self.assertEqual(lo.read(), "some")

        lo.truncate(6)
        lo.seek(0)
        # large object extended with zeroes
        self.assertEqual(lo.read(), "some\x00\x00")

        lo.truncate()
        lo.seek(0)
        # large object empty
        self.assertEqual(lo.read(), "")

    def test_truncate_after_close(self):
        if self.skip:
            return

        lo = self.conn.lobject()
        lo.close()
        self.assertRaises(psycopg2.InterfaceError, lo.truncate)

    def test_truncate_after_commit(self):
        if self.skip:
            return

        lo = self.conn.lobject()
        self.lo_oid = lo.oid
        self.conn.commit()

        self.assertRaises(psycopg2.ProgrammingError, lo.truncate)

def test_suite():
    return unittest.TestLoader().loadTestsFromName(__name__)

if __name__ == "__main__":
    unittest.main()