tests/lru_test.py

#!/usr/bin/env vpython3
# Copyright 2013 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

import json
import os
import tempfile
import unittest

# Mutates sys.path.
import test_env

from utils import lru


def _load_from_raw(state_text):
  """Makes a LRUDict by loading the given JSON from a file."""
  handle, tmp_name = tempfile.mkstemp(prefix='lru_test')
  os.close(handle)
  try:
    with open(tmp_name, 'w') as f:
      f.write(state_text)
    return lru.LRUDict.load(tmp_name)
  finally:
    try:
      os.unlink(tmp_name)
    except OSError:
      pass


def _save_and_load(lru_dict):
  """Saves then reloads a LRUDict instance."""
  handle, tmp_name = tempfile.mkstemp(prefix='lru_test')
  os.close(handle)
  try:
    lru_dict.save(tmp_name)
    return lru.LRUDict.load(tmp_name)
  finally:
    try:
      os.unlink(tmp_name)
    except OSError:
      pass


def _prepare_lru_dict(data):
  """Returns new LRUDict with given |keys| added one by one."""
  lru_dict = lru.LRUDict()
  for key, val in data:
    lru_dict.add(key, val)
  return lru_dict


class LRUDictTest(unittest.TestCase):
  def assert_same_data(self, expected, lru_dict):
    """Asserts that given |lru_dict| contains same data as |expected|.

    Tests items(), values(), get(), __iter__.
    """
    self.assertEqual(list(lru_dict.items()), expected)
    self.assertEqual(set(lru_dict), set(k for k, v in expected))
    self.assertEqual(list(lru_dict.values()), [v for k, v in expected])
    for k, v in expected:
      self.assertEqual(lru_dict.get(k), v)

  def test_empty(self):
    self.assert_same_data([], lru.LRUDict())

  def test_magic_methods_empty(self):
    """Tests __nonzero__, __iter, __len__, __getitem__ and __contains__."""
    # Check for empty dict.
    lru_dict = lru.LRUDict()
    self.assertFalse(lru_dict)
    self.assertEqual(len(lru_dict), 0)
    self.assertFalse(1 in lru_dict)
    self.assertFalse(list(lru_dict))
    with self.assertRaises(KeyError):
      _ = lru_dict[1]

  def test_magic_methods_nonempty(self):
    """Tests __nonzero__, __iter, __len__, __getitem__ and __contains__."""
    # Dict with one item.
    lru_dict = lru.LRUDict()
    lru_dict.add(1, 'one')
    self.assertTrue(lru_dict)
    self.assertEqual(len(lru_dict), 1)
    self.assertTrue(1 in lru_dict)
    self.assertFalse(2 in lru_dict)
    self.assertTrue(list(lru_dict))
    self.assertEqual('one', lru_dict[1])

  def test_add(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    lru_dict.add(1, 'one!!!')
    expected = [(2, 'two'), (3, 'three'), (1, 'one!!!')]
    self.assert_same_data(expected, lru_dict)
    lru_dict.add(0, 'zero')
    expected = [(2, 'two'), (3, 'three'), (1, 'one!!!'), (0, 'zero')]
    self.assert_same_data(expected, lru_dict)

  def test_pop_first(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    lru_dict.pop(1)
    self.assert_same_data([(2, 'two'), (3, 'three')], lru_dict)

  def test_pop_middle(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    lru_dict.pop(2)
    self.assert_same_data([(1, 'one'), (3, 'three')], lru_dict)

  def test_pop_last(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    lru_dict.pop(3)
    self.assert_same_data([(1, 'one'), (2, 'two')], lru_dict)

  def test_pop_missing(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    with self.assertRaises(KeyError):
      lru_dict.pop(4)

  def test_touch(self):
    lru_dict = _prepare_lru_dict([(1, 'one'), (2, 'two'), (3, 'three')])
    lru_dict.touch(2)
    self.assert_same_data([(1, 'one'), (3, 'three'), (2, 'two')], lru_dict)
    with self.assertRaises(KeyError):
      lru_dict.touch(4)

  def test_timestamp(self):
    """Tests get_oldest, pop_oldest."""
    lru_dict = lru.LRUDict()

    now = 0
    lru_dict.time_fn = lambda: now

    lru_dict.add('ka', 'va')
    now += 1

    lru_dict.add('kb', 'vb')
    now += 1

    self.assertEqual(lru_dict.get_oldest(), ('ka', ('va', 0)))
    self.assertEqual(lru_dict.pop_oldest(), ('ka', ('va', 0)))
    self.assertEqual(lru_dict.get_oldest(), ('kb', ('vb', 1)))
    self.assertEqual(lru_dict.pop_oldest(), ('kb', ('vb', 1)))

  def test_transform(self):
    lru_dict = lru.LRUDict()
    lru_dict.add('ka', 'va')
    lru_dict.add('kb', 'vb')
    lru_dict.transform(lambda k, v: v + '*')
    self.assert_same_data([('ka', 'va*'), ('kb', 'vb*')], lru_dict)

  def test_load_save_empty(self):
    self.assertFalse(_save_and_load(lru.LRUDict()))

  def test_load_save(self):
    data = [(1, None), (2, None), (3, None)]
    # Normal flow.
    lru_dict = _prepare_lru_dict(data)
    expected = [(1, None), (2, None), (3, None)]
    self.assert_same_data(expected, _save_and_load(lru_dict))

    # After touches.
    lru_dict = _prepare_lru_dict(data)
    lru_dict.touch(2)
    expected = [(1, None), (3, None), (2, None)]
    self.assert_same_data(expected, _save_and_load(lru_dict))

    # After pop.
    lru_dict = _prepare_lru_dict(data)
    lru_dict.pop(2)
    expected = [(1, None), (3, None)]
    self.assert_same_data(expected, _save_and_load(lru_dict))

    # After add.
    lru_dict = _prepare_lru_dict(data)
    lru_dict.add(4, 4)
    expected = [(1, None), (2, None), (3, None), (4, 4)]
    self.assert_same_data(expected, _save_and_load(lru_dict))

  def test_corrupted_state_file(self):
    # Loads correct state just fine.
    s = _load_from_raw(
        json.dumps({
            'version': lru.CURRENT_VERSION,
            'items': [
                ['key1', ['value1', 1]],
                ['key2', ['value2', 2]],
            ],
        }))
    self.assertIsNotNone(s)
    self.assertEqual(2, len(s))

    # Not a json.
    with self.assertRaises(ValueError):
      _load_from_raw('garbage, not a state')

    # Not a list.
    with self.assertRaises(ValueError):
      _load_from_raw('{}')

    # Not a list of pairs.
    with self.assertRaises(ValueError):
      _load_from_raw(json.dumps([
          ['key', 'value', 'and whats this?'],
      ]))

    # Duplicate keys.
    with self.assertRaises(ValueError):
      _load_from_raw(json.dumps([
          ['key', 'value'],
          ['key', 'another_value'],
      ]))


if __name__ == '__main__':
  test_env.main()