mindsdb

1
import datetime as dt
2
import time
3
import unittest
4
import traceback
5
import redis
6
import warnings
7
import tempfile
8
import json
9
import os
10

11
import pandas as pd
12

13
from mindsdb.utilities.cache import get_cache, RedisCache, FileCache, dataframe_checksum
14

15

16
class TestCashe(unittest.TestCase):
17

18
    @classmethod
19
    def setup_class(cls):
20
        # config
21
        config = {}
22
        # TODO run on own database
23
        fdi, cfg_file = tempfile.mkstemp(prefix='mindsdb_conf_')
24

25
        with os.fdopen(fdi, 'w') as fd:
26
            json.dump(config, fd)
27

28
        os.environ['MINDSDB_CONFIG_PATH'] = cfg_file
29

30
    def test_redis(self):
31
        cache = RedisCache('predict', max_size=2)
32
        try:
33
            self.cache_test(cache)
34
        except redis.ConnectionError as e:
35
            # Skip test for redis if no redis installed
36
            warnings.warn(f'redis is not available: {e}')
37
            print(traceback.format_exc())
38

39
    def test_file(self):
40
        cache = FileCache('predict', max_size=2)
41

42
        self.cache_test(cache)
43

44
    def cache_test(self, cache):
45

46
        # test save
47
        df = pd.DataFrame([
48
            [1, 1.2,'string', dt.datetime.now(), [1,2,3], {1:3}],
49
            [2, 3.2,'other', dt.datetime(2011, 12, 30), [3], {11:23, 2:3}],
50
        ], columns=['a', 'b', 'c', 'd', 'e', 'f'])
51

52
        # make bigger
53
        df = pd.concat([df]*100).reset_index()
54

55
        name = dataframe_checksum(df)
56

57
        # test save
58
        cache.set(name, df)
59

60
        df2 = cache.get(name)
61

62
        assert dataframe_checksum(df) == dataframe_checksum(df2)
63
        assert list(df.columns) == list(df2.columns)
64

65
        # test save df
66
        name += '1'
67
        cache.set_df(name, df)
68

69
        df2 = cache.get_df(name)
70

71
        assert dataframe_checksum(df) == dataframe_checksum(df2)
72
        assert list(df.columns) == list(df2.columns)
73

74
        # test delete
75
        cache.delete(name)
76

77
        df2 = cache.get(name)
78
        assert df2 is None
79

80
        # test max_size
81
        # load cache with size 2(max_size) + 5 (buffer)
82
        cache.set('first', df)
83
        for i in range(8):
84
            time.sleep(0.01)
85
            cache.set(str(i), df)
86

87
        # get first, must be deleted
88
        df2 = cache.get('first')
89
        assert df2 is None
90