From 3fb9beb5c04687ca9a266d8adf8589d830241a64 Mon Sep 17 00:00:00 2001
From: Stuart Gathman <stuart@gathman.org>
Date: Sat, 16 Feb 2013 05:40:46 +0000
Subject: [PATCH] Testcase for greylist

---
 Milter/greylist.py |  4 ++--
 test.py            |  2 ++
 testgrey.py        | 44 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 testgrey.py

diff --git a/Milter/greylist.py b/Milter/greylist.py
index 2ccf664..58d5dbc 100644
--- a/Milter/greylist.py
+++ b/Milter/greylist.py
@@ -35,7 +35,7 @@ class Greylist(object):
     self.dbp = shelve.open(dbname,'c',protocol=2)
     self.lock = thread.allocate_lock()
   
-  def check(self,ip,sender,recipient):
+  def check(self,ip,sender,recipient,timeinc=0):
     "Return number of allowed messages for greylist triple."
     sender = quoteAddress(sender)
     recipient = quoteAddress(recipient)
@@ -45,7 +45,7 @@ class Greylist(object):
       dbp = self.dbp
       try:
         r = dbp[key]
-        now = time.time()
+        now = time.time() + timeinc
         if now > r.lastseen + self.greylist_retain:
           # expired
           log.debug('Expired greylist: %s',key)
diff --git a/test.py b/test.py
index b52f367..895965e 100644
--- a/test.py
+++ b/test.py
@@ -2,6 +2,7 @@ import unittest
 import testmime
 import testsample
 import testutils
+import testgrey
 import os
 
 def suite(): 
@@ -9,6 +10,7 @@ def suite():
   s.addTest(testmime.suite())
   s.addTest(testsample.suite())
   s.addTest(testutils.suite())
+  s.addTest(testgrey.suite())
   return s
 
 if __name__ == '__main__':
diff --git a/testgrey.py b/testgrey.py
new file mode 100644
index 0000000..ee7fd0c
--- /dev/null
+++ b/testgrey.py
@@ -0,0 +1,44 @@
+import unittest
+import doctest
+import os
+import Milter.greylist
+from Milter.greylist import Greylist
+
+class GreylistTestCase(unittest.TestCase):
+
+  def setUp(self):
+    self.fname = 'test.db'
+
+  def tearDown(self):
+    os.remove(self.fname)
+
+  def testGrey(self):
+    grey = Greylist(self.fname)
+    # first time
+    rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com')
+    self.failUnless(rc == 0)
+    # not in window yet
+    rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=5*60)
+    self.failUnless(rc == 0)
+    # within window
+    rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=15*60)
+    self.failUnless(rc == 1)
+    # new triple
+    rc = grey.check('1.2.3.5','foo@bar.com','baz@spat.com',timeinc=15*60)
+    self.failUnless(rc == 0)
+    # seen again
+    rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=5*3600)
+    self.failUnless(rc == 2)
+    # new one past expire
+    rc = grey.check('1.2.3.5','foo@bar.com','baz@spat.com',timeinc=5*3600)
+    self.failUnless(rc == 0)
+    # original past retain 
+    rc = grey.check('1.2.3.4','foo@bar.com','baz@spat.com',timeinc=37*24*3600)
+    self.failUnless(rc == 0)
+
+def suite(): 
+  s = unittest.makeSuite(GreylistTestCase,'test')
+  return s
+
+if __name__ == '__main__':
+  unittest.TextTestRunner().run(suite())