From 8f58140419cbff8fb49c16c6c7e3d1317db6943f Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 3 Jan 2021 04:28:48 +0530
Subject: [PATCH] Start work on tests for disk cache

---
 kitty/disk-cache.c      | 35 ++++++++++++++++++++++++++++++++++-
 kitty/disk-cache.h      |  1 +
 kitty_tests/graphics.py | 27 ++++++++++++++++++++++++---
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/kitty/disk-cache.c b/kitty/disk-cache.c
index aabb402f5..e79fdbab6 100644
--- a/kitty/disk-cache.c
+++ b/kitty/disk-cache.c
@@ -20,6 +20,7 @@
 #include <stdlib.h>
 #include <sys/stat.h>
 #include <fcntl.h>
+#include <time.h>
 #ifdef HAS_SENDFILE
 #include <sys/sendfile.h>
 #endif
@@ -388,7 +389,7 @@ ensure_state(DiskCache *self) {
         PyObject *kc = NULL, *cache_dir = NULL;
         kc = PyImport_ImportModule("kitty.constants");
         if (kc) {
-            cache_dir = PyObject_CallMethod(kc, "dir_for_disk_cache", NULL);
+            cache_dir = PyObject_CallMethod(kc, "cache_dir", NULL);
             if (cache_dir) {
                 self->cache_dir = strdup(PyUnicode_AsUTF8(cache_dir));
                 if (!self->cache_dir) PyErr_NoMemory();
@@ -573,6 +574,29 @@ end:
     return data;
 }
 
+bool
+disk_cache_wait_for_write(PyObject *self_, monotonic_t timeout) {
+    DiskCache *self = (DiskCache*)self_;
+    monotonic_t end_at = monotonic() + timeout;
+    while (!timeout || monotonic() <= end_at) {
+        bool pending = false;
+        mutex(lock);
+        CacheEntry *s, *tmp;
+        HASH_ITER(hh, self->entries, s, tmp) {
+            if (!s->written_to_disk) {
+                pending = true;
+                break;
+            }
+        }
+        mutex(unlock);
+        if (!pending) return true;
+        wakeup_write_loop(self);
+        struct timespec a = { .tv_nsec = 50000000L}, b;
+        nanosleep(&a, &b);
+    }
+    return false;
+}
+
 #define PYWRAP(name) static PyObject* py##name(DiskCache *self, PyObject *args)
 #define PA(fmt, ...) if (!PyArg_ParseTuple(args, fmt, __VA_ARGS__)) return NULL;
 PYWRAP(ensure_state) {
@@ -594,6 +618,14 @@ PYWRAP(xor_data) {
     return ans;
 }
 
+static PyObject*
+wait_for_write(PyObject *self, PyObject *args) {
+    double timeout = 0;
+    PA("|d", &timeout);
+    if (disk_cache_wait_for_write(self, s_double_to_monotonic_t(timeout))) Py_RETURN_TRUE;
+    Py_RETURN_FALSE;
+}
+
 static PyObject*
 add(PyObject *self, PyObject *args) {
     const char *key, *data;
@@ -644,6 +676,7 @@ static PyMethodDef methods[] = {
     {"add", add, METH_VARARGS, NULL},
     {"remove", pyremove, METH_VARARGS, NULL},
     {"get", get, METH_VARARGS, NULL},
+    {"wait_for_write", wait_for_write, METH_VARARGS, NULL},
 
     {NULL}  /* Sentinel */
 };
diff --git a/kitty/disk-cache.h b/kitty/disk-cache.h
index de84334bf..17c1cbf7e 100644
--- a/kitty/disk-cache.h
+++ b/kitty/disk-cache.h
@@ -12,6 +12,7 @@ PyObject* create_disk_cache(void);
 bool add_to_disk_cache(PyObject *self, const void *key, size_t key_sz, const void *data, size_t data_sz);
 bool remove_from_disk_cache(PyObject *self_, const void *key, size_t key_sz);
 void* read_from_disk_cache(PyObject *self_, const void *key, size_t key_sz, void*(allocator)(void*, size_t), void*);
+bool disk_cache_wait_for_write(PyObject *self, monotonic_t timeout);
 
 static inline void* disk_cache_malloc_allocator(void *x, size_t sz) {
     *((size_t*)x) = sz;
diff --git a/kitty_tests/graphics.py b/kitty_tests/graphics.py
index c8d82c440..85a4f6f74 100644
--- a/kitty_tests/graphics.py
+++ b/kitty_tests/graphics.py
@@ -3,7 +3,6 @@
 # License: GPL v3 Copyright: 2016, Kovid Goyal <kovid at kovidgoyal.net>
 
 import os
-import shutil
 import tempfile
 import unittest
 import zlib
@@ -138,10 +137,9 @@ class TestGraphics(BaseTest):
 
     def setUp(self):
         self.cache_dir = cache_dir.override_dir = tempfile.mkdtemp()
-        self.disk_cache_dir = os.path.join(self.cache_dir, 'disk-cache')
 
     def tearDown(self):
-        shutil.rmtree(self.cache_dir)
+        os.rmdir(self.cache_dir)
         cache_dir.override_dir = None
 
     def test_xor_data(self):
@@ -157,6 +155,29 @@ class TestGraphics(BaseTest):
                 data = base + base_data[:extra]
                 self.assertEqual(xor_data(key, data), xor(key, data))
 
+    def test_disk_cache(self):
+        s = self.create_screen()
+        dc = s.grman.disk_cache
+        data = {}
+
+        def key_as_bytes(key):
+            if isinstance(key, int):
+                key = str(key)
+            if isinstance(key, str):
+                key = key.encode('utf-8')
+            return bytes(key)
+
+        def add(key, val):
+            bkey = key_as_bytes(key)
+            data[key] = key_as_bytes(val)
+            dc.add(bkey, data[key])
+
+        for i in range(25):
+            self.assertIsNone(add(i, f'{i}' * i))
+
+        self.assertEqual(dc.total_size, sum(map(len, data.values())))
+        self.assertTrue(dc.wait_for_write())
+
     def test_load_images(self):
         s, g, l, sl = load_helpers(self)