forked from Jelloeater/NetworkMonitor
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_copy.py
108 lines (95 loc) · 3.61 KB
/
test_copy.py
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
import unittest
import pg8000
from .connection_settings import db_connect
from pg8000.six import b, BytesIO
from sys import exc_info
class Tests(unittest.TestCase):
def setUp(self):
self.db = pg8000.connect(**db_connect)
try:
cursor = self.db.cursor()
try:
cursor = self.db.cursor()
cursor.execute("DROP TABLE t1")
except pg8000.DatabaseError:
e = exc_info()[1]
# the only acceptable error is:
self.assertEqual(
e.args[1], b('42P01'), # table does not exist
"incorrect error for drop table")
self.db.rollback()
cursor.execute(
"CREATE TEMPORARY TABLE t1 (f1 int primary key, "
"f2 int not null, f3 varchar(50) null)")
finally:
cursor.close()
def tearDown(self):
self.db.close()
def testCopyToWithTable(self):
try:
cursor = self.db.cursor()
cursor.execute(
"INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (1, 1, 1))
cursor.execute(
"INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (2, 2, 2))
cursor.execute(
"INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (3, 3, 3))
stream = BytesIO()
cursor.copy_to(stream, "t1")
self.assertEqual(
stream.getvalue(), b("1\t1\t1\n2\t2\t2\n3\t3\t3\n"))
self.assertEqual(cursor.rowcount, 3)
self.db.commit()
finally:
cursor.close()
def testCopyToWithQuery(self):
try:
cursor = self.db.cursor()
stream = BytesIO()
cursor.copy_to(
stream, query="COPY (SELECT 1 as One, 2 as Two) TO STDOUT "
"WITH DELIMITER 'X' CSV HEADER QUOTE AS 'Y' FORCE QUOTE Two")
self.assertEqual(stream.getvalue(), b('oneXtwo\n1XY2Y\n'))
self.assertEqual(cursor.rowcount, 1)
self.db.rollback()
finally:
cursor.close()
def testCopyFromWithTable(self):
try:
cursor = self.db.cursor()
stream = BytesIO(b("1\t1\t1\n2\t2\t2\n3\t3\t3\n"))
cursor.copy_from(stream, "t1")
self.assertEqual(cursor.rowcount, 3)
cursor.execute("SELECT * FROM t1 ORDER BY f1")
retval = cursor.fetchall()
self.assertEqual(retval, ([1, 1, '1'], [2, 2, '2'], [3, 3, '3']))
self.db.rollback()
finally:
cursor.close()
def testCopyFromWithQuery(self):
try:
cursor = self.db.cursor()
stream = BytesIO(b("f1Xf2\n1XY1Y\n"))
cursor.copy_from(
stream, query="COPY t1 (f1, f2) FROM STDIN WITH DELIMITER "
"'X' CSV HEADER QUOTE AS 'Y' FORCE NOT NULL f1")
self.assertEqual(cursor.rowcount, 1)
cursor.execute("SELECT * FROM t1 ORDER BY f1")
retval = cursor.fetchall()
self.assertEqual(retval, ([1, 1, None],))
self.db.commit()
finally:
cursor.close()
def testCopyWithoutTableOrQuery(self):
try:
cursor = self.db.cursor()
stream = BytesIO()
self.assertRaises(
pg8000.CopyQueryOrTableRequiredError, cursor.copy_from, stream)
self.assertRaises(
pg8000.CopyQueryOrTableRequiredError, cursor.copy_to, stream)
self.db.rollback()
finally:
cursor.close()
if __name__ == "__main__":
unittest.main()