outlines

1
import os
2
import tempfile
3

4
import diskcache
5
import pytest
6

7

8
@pytest.fixture
9
def refresh_environment():
10
    """Refresh the test environment.
11

12
    This deletes any reference to `outlines` in the modules dictionary and unsets the
13
    `OUTLINES_CACHE_DIR` environment variable if set. This is necessary because we
14
    are using a module variable to hold the cache.
15

16
    """
17
    import sys
18

19
    for key in list(sys.modules.keys()):
20
        if "outlines" in key:
21
            del sys.modules[key]
22

23
    try:
24
        del os.environ["OUTLINES_CACHE_DIR"]
25
    except KeyError:
26
        pass
27

28

29
@pytest.fixture
30
def test_cache(refresh_environment):
31
    """Initialize a temporary cache and delete it after the test has run."""
32
    with tempfile.TemporaryDirectory() as tempdir:
33
        os.environ["OUTLINES_CACHE_DIR"] = tempdir
34
        import outlines
35

36
        memory = outlines.get_cache()
37
        assert memory.directory == tempdir
38

39
        yield outlines.caching.cache()
40

41
        memory.clear()
42

43

44
def test_get_cache(test_cache):
45
    import outlines
46

47
    memory = outlines.get_cache()
48
    assert isinstance(memory, diskcache.Cache)
49

50
    # If the cache is enabled then the size
51
    # of `store` should not increase the
52
    # second time `f` is called.
53
    store = list()
54

55
    @test_cache
56
    def f(x):
57
        store.append(1)
58
        return x
59

60
    f(1)
61
    store_size = len(store)
62

63
    f(1)
64
    assert len(store) == store_size
65

66
    f(2)
67
    assert len(store) == store_size + 1
68

69

70
def test_disable_cache(test_cache):
71
    """Make sure that we can disable the cache."""
72
    import outlines
73

74
    outlines.disable_cache()
75

76
    # If the cache is disabled then the size
77
    # of `store` should increase every time
78
    # `f` is called.
79
    store = list()
80

81
    @test_cache
82
    def f(x):
83
        store.append(1)
84
        return x
85

86
    f(1)
87
    store_size = len(store)
88
    f(1)
89
    assert len(store) == store_size + 1
90

91

92
def test_clear_cache(test_cache):
93
    """Make sure that we can clear the cache."""
94
    import outlines
95

96
    store = list()
97

98
    @test_cache
99
    def f(x):
100
        store.append(1)
101
        return x
102

103
    # The size of `store` does not increase since
104
    # `f` is cached after the first run.
105
    f(1)
106
    store_size = len(store)
107
    f(1)
108
    assert len(store) == store_size
109

110
    # The size of `store` should increase if we call `f`
111
    # after clearing the cache.
112
    outlines.clear_cache()
113
    f(1)
114
    assert len(store) == store_size + 1
115

116

117
def test_version_upgrade_cache_invalidate(test_cache, mocker):
118
    """Ensure we can change the signature of a cached function if we upgrade the version"""
119

120
    import outlines.caching
121

122
    def simulate_restart_outlines():
123
        # clearing in-memory lru_cache which returns the diskcache in
124
        # order to simulate a reload, we're not clearing the diskcache itself
125
        outlines.caching.get_cache.cache_clear()
126

127
    mocker.patch("outlines._version.__version__", new="0.0.0")
128
    simulate_restart_outlines()
129

130
    # initialize cache with signature of Tuple-of-3
131
    @test_cache
132
    def foo():
133
        return (1, 2, 3)
134

135
    a, b, c = foo()
136

137
    # "restart" outlines without upgrading version
138
    simulate_restart_outlines()
139

140
    # change signature to Tuple-of-2
141
    @test_cache
142
    def foo():
143
        return (1, 2)
144

145
    # assert without version upgrade, old, bad cache is used
146
    with pytest.raises(ValueError):
147
        a, b = foo()
148

149
    # "restart" outlines WITH version upgrade
150
    mocker.patch("outlines._version.__version__", new="0.0.1")
151
    simulate_restart_outlines()
152

153
    # change signature to Tuple-of-2
154
    @test_cache
155
    def foo():
156
        return (1, 2)
157

158
    # assert with version upgrade, old cache is invalidated and new cache is used
159
    a, b = foo()
160