CoCalc -- test_results.py

GitHub Repository: singlestore-labs/singlestoredb-python
Path: blob/main/singlestoredb/tests/test_results.py
⁴⁶⁹ views
1
#!/usr/bin/env python
2
# type: ignore
3
"""Test SingleStoreDB results."""
4
import os
5
import unittest
6

7
import pandas as pd
8

9
import singlestoredb as s2
10
from singlestoredb.tests import utils
11
# import traceback
12

13

14
class TestResults(unittest.TestCase):
15

16
    dbname: str = ''
17
    dbexisted: bool = False
18

19
    @classmethod
20
    def setUpClass(cls):
21
        sql_file = os.path.join(os.path.dirname(__file__), 'test.sql')
22
        cls.dbname, cls.dbexisted = utils.load_sql(sql_file)
23

24
    @classmethod
25
    def tearDownClass(cls):
26
        if not cls.dbexisted:
27
            utils.drop_database(cls.dbname)
28

29
    def setUp(self):
30
        self.conn = s2.connect(database=type(self).dbname)
31
        self.cur = self.conn.cursor()
32

33
    def tearDown(self):
34
        try:
35
            if self.cur is not None:
36
                self.cur.close()
37
        except Exception:
38
            # traceback.print_exc()
39
            pass
40

41
        try:
42
            if self.conn is not None:
43
                self.conn.close()
44
        except Exception:
45
            # traceback.print_exc()
46
            pass
47

48
    def test_tuples(self):
49
        with s2.options(('results.type', 'tuples')):
50
            with s2.connect(database=type(self).dbname) as conn:
51
                with conn.cursor() as cur:
52
                    cur.execute('select * from data')
53
                    out = cur.fetchone()
54
                    assert type(out) is tuple, type(out)
55
                    assert len(out) == 3, len(out)
56
                    cur.fetchall()
57

58
                    cur.execute('select * from data')
59
                    out = cur.fetchall()
60
                    assert len(out) == 5, len(out)
61
                    assert len(out[0]) == 3, len(out[0])
62
                    assert type(out[0]) is tuple, type(out[0])
63
                    assert sorted(out) == sorted([
64
                        ('a', 'antelopes', 2),
65
                        ('b', 'bears', 2),
66
                        ('c', 'cats', 5),
67
                        ('d', 'dogs', 4),
68
                        ('e', 'elephants', 0),
69
                    ]), out
70

71
                    out = cur.fetchall()
72
                    assert len(out) == 0, len(out)
73

74
    def test_namedtuples(self):
75
        with s2.options(('results.type', 'namedtuples')):
76
            with s2.connect(database=type(self).dbname) as conn:
77
                with conn.cursor() as cur:
78
                    cur.execute('select * from data')
79
                    out = cur.fetchone()
80
                    assert isinstance(out, tuple), type(out)
81
                    assert len(out) == 3, len(out)
82
                    assert type(out).__name__ == 'Row', type(out).__name__
83
                    assert hasattr(out, 'id')
84
                    assert hasattr(out, 'name')
85
                    assert hasattr(out, 'value')
86
                    cur.fetchall()
87

88
                    cur.execute('select * from data')
89
                    out = cur.fetchall()
90
                    assert len(out) == 5, len(out)
91
                    assert len(out[0]) == 3, len(out[0])
92
                    assert isinstance(out[0], tuple), type(out[0])
93
                    assert type(out[0]).__name__ == 'Row', type(out[0]).__name__
94
                    assert hasattr(out[0], 'id')
95
                    assert hasattr(out[0], 'name')
96
                    assert hasattr(out[0], 'value')
97
                    assert sorted(out) == sorted([
98
                        ('a', 'antelopes', 2),
99
                        ('b', 'bears', 2),
100
                        ('c', 'cats', 5),
101
                        ('d', 'dogs', 4),
102
                        ('e', 'elephants', 0),
103
                    ]), out
104

105
                    out = cur.fetchall()
106
                    assert len(out) == 0, len(out)
107

108
    def test_dict(self):
109
        with s2.options(('results.type', 'dicts')):
110
            with s2.connect(database=type(self).dbname) as conn:
111
                with conn.cursor() as cur:
112
                    cur.execute('select * from data')
113
                    out = cur.fetchone()
114
                    assert type(out) is dict, type(out)
115
                    assert len(out) == 3, len(out)
116
                    cur.fetchall()
117

118
                    cur.execute('select * from data')
119
                    out = cur.fetchall()
120
                    assert type(out[0]) is dict, type(out[0])
121
                    assert len(out) == 5, len(out)
122
                    assert len(out[0]) == 3, len(out[0])
123
                    assert sorted(out, key=lambda x: x['id']) == sorted(
124
                        [
125
                            dict(id='a', name='antelopes', value=2),
126
                            dict(id='b', name='bears', value=2),
127
                            dict(id='c', name='cats', value=5),
128
                            dict(id='d', name='dogs', value=4),
129
                            dict(id='e', name='elephants', value=0),
130
                        ], key=lambda x: x['id'],
131
                    )
132

133
                    out = cur.fetchall()
134
                    assert len(out) == 0, len(out)
135

136
    def _test_dataframe(self):
137
        with s2.options(('results.type', 'dataframe')):
138
            with s2.connect(database=type(self).dbname) as conn:
139
                with conn.cursor() as cur:
140
                    cur.execute('select * from data')
141
                    out = cur.fetchone()
142
                    assert type(out) is pd.DataFrame, type(out)
143
                    assert len(out) == 1, len(out)
144
                    cur.fetchall()
145

146
                    cur.execute('select * from data')
147
                    out = cur.fetchall()
148
                    assert type(out) is pd.DataFrame, type(out)
149
                    assert len(out) == 5, len(out)
150
                    out = out.sort_values('id').reset_index(drop=True)
151
                    exp = pd.DataFrame(
152
                        data=[
153
                            ('a', 'antelopes', 2),
154
                            ('b', 'bears', 2),
155
                            ('c', 'cats', 5),
156
                            ('d', 'dogs', 4),
157
                            ('e', 'elephants', 0),
158
                        ], columns=['id', 'name', 'value'],
159
                    ).sort_values('id').reset_index(drop=True)
160
                    assert list(out.columns) == list(exp.columns), list(out.columns)
161
                    assert [list(x) for x in list(out.values)] == \
162
                           [list(x) for x in list(exp.values)], \
163
                           [list(x) for x in list(out.values)]
164

165
                    out = cur.fetchall()
166
                    assert len(out) == 0, len(out)
167

168

169
if __name__ == '__main__':
170
    import nose2
171
    nose2.main()
172

173
Product

Resources

Company
