pytorch

1

2

3

4

5

6
from caffe2.python import workspace
7

8
import os
9
import tempfile
10
import unittest
11

12

13
class TestDB(unittest.TestCase):
14
    def setUp(self):
15
        handle, self.file_name = tempfile.mkstemp()
16
        os.close(handle)
17
        self.data = [
18
            (
19
                "key{}".format(i).encode("ascii"),
20
                "value{}".format(i).encode("ascii")
21
            )
22
            for i in range(1, 10)
23
        ]
24

25
    def testSimple(self):
26
        db = workspace.C.create_db(
27
            "minidb", self.file_name, workspace.C.Mode.write)
28

29
        for key, value in self.data:
30
            transaction = db.new_transaction()
31
            transaction.put(key, value)
32
            del transaction
33

34
        del db  # should close DB
35

36
        db = workspace.C.create_db(
37
            "minidb", self.file_name, workspace.C.Mode.read)
38
        cursor = db.new_cursor()
39
        data = []
40
        while cursor.valid():
41
            data.append((cursor.key(), cursor.value()))
42
            cursor.next()  # noqa: B305
43
        del cursor
44

45
        db.close()  # test explicit db closer
46
        self.assertEqual(data, self.data)
47