modules/predict: added test for prediction process

this tests that:
* sampling frequent queries works
* the code to find periodic appearences of the same name and type
* resolving predicted queries works

it doesn't test pessimistic cases or failure modes

diff --git a/modules/predict/predict.lua b/modules/predict/predict.lua
index d0f550ba84f5ff8d7eb8c8c8fac5a5fd93c0d367..5160b03fe4a10a48923fe745d1157958e7a3ab47 100644 (file)
--- a/modules/predict/predict.lua
+++ b/modules/predict/predict.lua
@@ -16,19 +16,19 @@ local predict = {
 -- Load dependent modules
 if not stats then modules.load('stats') end
 
--- Calculate current epoch (which window fits current time)
-local function current_epoch()
-       if not predict.period or predict.period <= 1 then return nil end
-       return (os.date('%H')*(60/predict.window) +
-               math.floor(os.date('%M')/predict.window)) % predict.period + 1
-end
-
 -- Calculate next sample with jitter [1-2/5 of window]
 local function next_event()
        local jitter = (predict.window * minute) / 5;
        return math.random(jitter, 2 * jitter)
 end
 
+-- Calculate current epoch (which window fits current time)
+function predict.epoch()
+       if not predict.period or predict.period <= 1 then return nil end
+       return (os.date('%H')*(60/predict.window) +
+               math.floor(os.date('%M')/predict.window)) % predict.period + 1
+end
+
 -- Resolve queued records and flush the queue
 function predict.drain()
        local deleted = 0
@@ -37,10 +37,13 @@ function predict.drain()
                worker.resolve(qname, kres.type[qtype], kres.class.IN, 'NO_CACHE')
                predict.queue[key] = nil
                deleted = deleted + 1
-               if deleted >= predict.batch then
+               -- Resolve smaller batches at a time
+               if predict.batch > 0 and deleted >= predict.batch then
                        break
                end
        end
+       -- Schedule prefetch of another batch if not complete
+       if predict.ev_drain then event.cancel(predict.ev_drain) end
        predict.ev_drain = nil
        if deleted > 0 then
                predict.ev_drain = event.after((predict.window * 3) * sec, predict.drain)
@@ -101,14 +104,13 @@ end
 
 function predict.process()
        -- Start a new epoch, or continue sampling
-       predict.ev_sample = nil
-       local epoch_now = current_epoch()
+       local epoch_now = predict.epoch()
        local nr_queued = 0
 
        -- End of epoch
-       if predict.epoch ~= epoch_now then
+       if predict.current_epoch ~= epoch_now then
                stats['predict.epoch'] = epoch_now
-               predict.epoch = epoch_now
+               predict.current_epoch = epoch_now
                -- enqueue records from upcoming epoch
                nr_queued = enqueue_from_log(predict.log[epoch_now])
                -- predict next epoch
@@ -128,6 +130,8 @@ function predict.process()
                        predict.ev_drain = event.after(0, predict.drain)
                end
        end
+
+       if predict.ev_sample then event.cancel(predict.ev_sample) end
        predict.ev_sample = event.after(next_event(), predict.process)
        if stats then
                stats['predict.queue'] = predict.queue_len
@@ -138,7 +142,7 @@ end
 
 function predict.init()
        if predict.window > 0 then
-               predict.epoch = current_epoch()
+               predict.current_epoch = predict.epoch()
                predict.ev_sample = event.after(next_event(), predict.process)
        end
 end

diff --git a/tests/config/predict/test.cfg b/tests/config/predict/test.cfg
index c0c1069baa7539b0d4e764fd99658715c89a2136..3344f97bd799ddc7036edf1da39f18d8f903a64d 100644 (file)
--- a/tests/config/predict/test.cfg
+++ b/tests/config/predict/test.cfg
@@ -3,16 +3,58 @@ dofile('./test_utils.lua') -- load test utilities
 -- setup resolver
 modules = { 'predict' }
 
--- test if prediction of non-standard types works
-function test_predict_drain_typex()
-       predict.queue_len = 1
+-- mock global functions
+local resolve_count = 0
+worker.resolve = function ()
+       resolve_count = resolve_count + 1
+end
+stats.frequent = function ()
+       return {
+               {name = 'example.com', type = 'TYPE65535'},
+               {name = 'example.com', type = 'SOA'},
+       }
+end
+local current_epoch = 0
+predict.epoch = function ()
+       return current_epoch % predict.period + 1
+end
+
+-- test if draining of prefetch queue works
+function test_predict_drain()
+       predict.queue_len = 2
        predict.queue['TYPE65535 example.com'] = 1
+       predict.queue['SOA example.com'] = 1
+       predict.drain()
+       -- test that it attempted to prefetch
+       assert.same(2, resolve_count)
+       assert.same(0, predict.queue_len)
+end
+
+-- test if prediction process works
+function test_predict_process()
+       -- start new epoch
+       predict.process()
+       assert.same(0, predict.queue_len)
+       -- next epoch, still no period for frequent queries
+       current_epoch = current_epoch + 1
+       predict.process()
+       assert.same(0, predict.queue_len)
+       -- next epoch, found period
+       current_epoch = current_epoch + 1
+       predict.process()
+       assert.same(2, predict.queue_len)
+       -- drain works with scheduled prefetches (two batches)
+       resolve_count = 0
+       predict.drain()
        predict.drain()
+       assert.same(2, resolve_count)
+       assert.same(0, predict.queue_len)
 end
 
 -- run test after processed config file
 -- default config will be used and we can test it.
 event.after(0, function (ev) 
-       test(test_predict_drain_typex)
+       test(test_predict_drain)
+       test(test_predict_process)
        quit()
 end)