1import pytest
2import os
3import tempfile
4import time
5import six
6
7from dlmanager import fs
8from dlmanager.persist_limit import PersistLimit
9
10
11class TempCreator(object):
12    def __init__(self):
13        self.tempdir = tempfile.mkdtemp()
14
15    def list(self):
16        return os.listdir(self.tempdir)
17
18    def create_file(self, name, size, delay):
19        fname = os.path.join(self.tempdir, name)
20        with open(fname, 'wb') as f:
21            f.write(six.b('a' * size))
22        # equivalent to touch, but we apply a delay for the test
23        atime = time.time() + delay
24        os.utime(fname, (atime, atime))
25
26
27@pytest.yield_fixture
28def temp():
29    tmp = TempCreator()
30    yield tmp
31    fs.remove(tmp.tempdir)
32
33
34@pytest.mark.parametrize("size_limit,file_limit,files", [
35    # limit_file is always respected
36    (10, 5, "bcdef"),
37    (10, 3, "def"),
38    # if size_limit or file_limit is 0, nothing is removed
39    (0, 5, "abcdef"),
40    (5, 0, "abcdef"),
41    # limit_size works
42    (35, 1, "def"),
43])
44def test_persist_limit(temp, size_limit, file_limit, files):
45    temp.create_file("a", 10, -6)
46    temp.create_file("b", 10, -5)
47    temp.create_file("c", 10, -4)
48    temp.create_file("d", 10, -3)
49    temp.create_file("e", 10, -2)
50    temp.create_file("f", 10, -1)
51
52    persist_limit = PersistLimit(size_limit, file_limit)
53    persist_limit.register_dir_content(temp.tempdir)
54    persist_limit.remove_old_files()
55
56    assert ''.join(sorted(temp.list())) == ''.join(sorted(files))
57