#!/usr/bin/python

import unittest
import sys
import time

sys.path.insert(0, '../src/lib')
import LookupManager

class LookupTestCase(unittest.TestCase):
    def setUp(self):
        self.lm = LookupManager.get_instance()
        self.timeout = 60
        self.extraNames = (
            "art.gnome.org", "b9.com", "blog.neoteny.com",
            "diveintomark.org", "freshmeat.net", "home.comcast.net",
            "imperialviolet.org", "inessential.com", "kode-fu.com",
            "kt.zork.net", "lambda.weblogs.com",
            "morenews.blogspot.com", "mpt.phrasewise.com",
            "online.effbot.org", "patricklogan.blogspot.com",
            "people.enginesofcreation.ie", "polydistortion.net",
            "pragprog.com", "primates.ximian.com",
            "radio.weblogs.com", "rover.cs.northwestern.edu",
            "slashdot.org", "tbray.org", "toulouse.amber.org",
            "usefulinc.com", "weblog.infoworld.com",
            "weblogs.mozillazine.org", "wmf.editthispage.com",
            "ww.telent.net", "www.advogato.org", "www.amk.ca",
            "www.artima.com", "www.benhammersley.com",
            "www.blackbeltjones.com", "www.burtonini.com",
            "www.deadlybloodyserious.com", "www.debianplanet.net",
            "www.gnomedesktop.org", "www.hole.fi", "www.iki.fi",
            "www.inkstain.net", "www.lemonodor.com",
            "www.markokarppinen.com", "www.markwatson.com",
            "www.merrells.com", "www.mozillazine.org", "www.nat.org",
            "www.nongnu.org", "www.orbtech.com",
            "www.playbynumbers.com", "www.python.org",
            "www.rekaksois.com", "www.s2.org", "www.sidhe.org",
            "www.spacetoday.net", "www.sweetcode.org",
            "www.textism.com", "www.toolshed.com", "www.zanthan.com")

    def _callback(self, name, ip, data=None):
        self.response = (name, ip)

    def testGoodARecord(self):
        self.response = None
        self.lm.lookup('www.google.com', self._callback)
        now = time.time()
        while not self.response:
            assert time.time() - now < self.timeout, 'Lookup timed out'
            self.lm.poll(1)
        assert self.response[1] is not None

    def testBadARecord(self):
        self.response = None
        # this has, at least at some point, returned a cname even though
        # asked for an a record
        self.lm.lookup('www.morestuff4less.com', self._callback)
        now = time.time()
        while not self.response:
            assert time.time() - now < self.timeout, 'Lookup timed out'
            self.lm.poll(1)
        assert self.response[1] is not None

    def testEmptyDomain(self):
        self.response = None
        self.lm.lookup('', self._callback)
        now = time.time()
        while not self.response:
            assert time.time() - now < self.timeout, 'Lookup timed out'
            self.lm.poll(1)
        assert self.response[1] is None

    def testInvalidDomain(self):
        self.response = None
        self.lm.lookup('www.sdkfskdjjfsd.xcv', self._callback)
        now = time.time()
        while not self.response:
            assert time.time() - now < self.timeout, 'Lookup timed out'
            self.lm.poll(1)
        assert self.response[1] is None

    def _callback_multiple(self, name, ip, data=None):
        print "received %s for %s" % (name, ip)
        self.multiple_responses[name] = (name, ip)

    def testExtraNames(self):
        self.multiple_responses = {}
        for n in self.extraNames:
            self.multiple_responses[n] = None
            self.lm.lookup(n, self._callback_multiple)
        now = time.time()
        timeout = self.timeout * len(self.extraNames)
        while None in self.multiple_responses.values():
            assert time.time() - now < self.timeout, 'Lookup timed out'
            self.lm.poll(1)
        assert None not in self.multiple_responses.values()

def suite():
    suite = unittest.TestSuite()
    suite.addTest(LookupTestCase('testGoodARecord'))
    suite.addTest(LookupTestCase('testBadARecord'))
    suite.addTest(LookupTestCase('testEmptyDomain'))
    suite.addTest(LookupTestCase('testInvalidDomain'))
    suite.addTest(LookupTestCase('testExtraNames'))
    return suite

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