[tor-commits] [nyx/master] Add cache helper

[atagar at torproject.org]

commit e9fbe44319256790aa69e8e20ebec042511a55d8
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Aug 28 12:12:15 2017 -0700

    Add cache helper
    
    Simple helper that provides back a sqlite cache for us to use. If we can
    persist it then great, if not then it resides in memory.
---
 nyx/__init__.py | 18 ++++++++++++++++++
 test/cache.py   | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/nyx/__init__.py b/nyx/__init__.py
index 3cf3c85..c7db0cb 100644
--- a/nyx/__init__.py
+++ b/nyx/__init__.py
@@ -33,6 +33,7 @@ Tor curses monitoring application.
 
 import distutils.spawn
 import os
+import sqlite3
 import sys
 import threading
 import time
@@ -84,6 +85,7 @@ CONFIG = stem.util.conf.config_dict('nyx', {
 NYX_INTERFACE = None
 TOR_CONTROLLER = None
 BASE_DIR = os.path.sep.join(__file__.split(os.path.sep)[:-1])
+CACHE = None
 
 # technically can change but we use this query a *lot* so needs to be cached
 
@@ -250,6 +252,22 @@ def data_directory(filename, config):
   return os.path.join(data_dir, filename)
 
 
+def cache():
+  """
+  Provides the sqlite cache for application data.
+
+  :returns: **sqlite3.Connection** for our applicaion cache
+  """
+
+  global CACHE
+
+  if CACHE is None:
+    cache_path = data_directory('cache.sqlite')
+    CACHE = sqlite3.connect(cache_path if cache_path else ':memory:')
+
+  return CACHE
+
+
 @uses_settings
 def expand_path(path, config):
   """
diff --git a/test/cache.py b/test/cache.py
new file mode 100644
index 0000000..dc1fa40
--- /dev/null
+++ b/test/cache.py
@@ -0,0 +1,40 @@
+"""
+Unit tests for nyx.cache.
+"""
+
+import tempfile
+import unittest
+
+import nyx
+
+from mock import Mock, patch
+
+
+class TestCache(unittest.TestCase):
+  def setUp(self):
+    nyx.CACHE = None  # drop cached database reference
+
+  @patch('nyx.data_directory', Mock(return_value = None))
+  def test_memory_cache(self):
+    cache = nyx.cache()
+    self.assertEqual((0, 'main', ''), cache.execute("PRAGMA database_list").fetchone())
+
+    cache.execute('CREATE TABLE aliases(alias TEXT, command TEXT)')
+    cache.execute('INSERT INTO aliases(alias, command) VALUES (?,?)', ('l', 'ls -xF --color=auto'))
+    cache.execute('INSERT INTO aliases(alias, command) VALUES (?,?)', ('ll', 'ls -hlA --color=auto'))
+    self.assertEqual('ls -hlA --color=auto', cache.execute('SELECT command FROM aliases WHERE alias=?', ('ll',)).fetchone()[0])
+
+  def test_file_cache(self):
+    with tempfile.NamedTemporaryFile(suffix = '.sqlite') as tmp:
+      with patch('nyx.data_directory', Mock(return_value = tmp.name)):
+        cache = nyx.cache()
+        self.assertEqual((0, 'main', tmp.name), cache.execute("PRAGMA database_list").fetchone())
+
+        cache.execute('CREATE TABLE aliases(alias TEXT, command TEXT)')
+        cache.execute('INSERT INTO aliases(alias, command) VALUES (?,?)', ('l', 'ls -xF --color=auto'))
+        cache.execute('INSERT INTO aliases(alias, command) VALUES (?,?)', ('ll', 'ls -hlA --color=auto'))
+        cache.commit()
+        cache.close()
+        nyx.CACHE = None
+
+        self.assertEqual('ls -hlA --color=auto', nyx.cache().execute('SELECT command FROM aliases WHERE alias=?', ('ll',)).fetchone()[0])