#!/usr/bin/env python

import unittest, time, string, StringIO
from nws.client import *

try:
    from logging import getLogger, basicConfig, DEBUG
except ImportError:
    DEBUG = 1
    class Logger:
        def __init__(self):
            self._level = 0

        def info(self, msg):
            if self._level:
                print msg

        def setLevel(self, level):
            self._level = level

    def getLogger(ign):
        return Logger()

    def basicConfig():
        pass

def getdict(**kw):
    return kw

class NWSTest(unittest.TestCase):
    log = getLogger('NWSTest')
    wsname = 'unit test ws'

    def setUp(self):
        self.log.info('Setting up')
        self.ws = NetWorkSpace(self.wsname, **nwsargs)

    def tearDown(self):
        self.log.info('Tearing down')
        try: self.ws.server.close()
        except: pass

    # Get a list of all workspaces from a listing string generated by
    # the listWss method.
    def getWss(self, listing):
        x = {}
        self.assertEqual(listing[-1], '\n')
        listing = listing[:-1]
        for line in listing.split('\n'):
            t = line.split('\t')
            self.assertEqual(len(t), 5)
            self.assert_(len(t[0]) > 1)
            if t[0][0] == '>':
                mine = 1
            else:
                self.assertEqual(t[0][0], ' ')
                mine = 0
            name = t[0][1:]
            owner = t[1]
            if t[2] == 'True':
                persistent = 1
            else:
                self.assertEqual(t[2], 'False')
                persistent = 0
            nvars = int(t[3])
            self.assert_(nvars >= 0)
            if t[4]:
                vars = t[4].split(',')
            else:
                vars = []
            self.assertEqual(nvars, len(vars))
            x[name] = getdict(name=name, mine=mine, owner=owner,
                    persistent=persistent, nvars=nvars, vars=vars)

        return x

    # Get a list of all variables from a listing string generated by
    # the listVars method.
    def getVars(self, listing):
        x = {}
        if listing:
            self.assert_(listing[-1] != '\n')
            lines = listing.split('\n')
            for line in lines:
                t = line.split('\t')
                self.assertEqual(len(t), 5)
                self.assert_(len(t[0]) > 0)
                variable = t[0]
                values = int(t[1])
                fetchers = int(t[2])
                finders = int(t[3])
                mode = t[4]
                x[variable] = getdict(variable=variable, values=values,
                        fetchers=fetchers, finders=finders, mode=mode)

        return x

    def getVarsDict(self, varDict):
        x = {}
        for k in varDict:
            v = varDict[k]
            x[v[V_VARIABLE]] = getdict(variable=v[V_VARIABLE], values=v[V_VALUES],
                    fetchers=v[V_FETCHERS], finders=v[V_FINDERS], mode=v[V_MODE])

        return x

    def testFifo(self):
        """A basic test of using a fifo mode variable"""
        values = range(10)

        self.log.info('Storing variables in FIFO mode')
        for i in values:
            self.ws.store('FIFO', i)

        self.log.info('Fetching values for FIFO variable')
        for i in values:
            t = self.ws.fetchTry('FIFO')
            self.assertEqual(t, i)

        self.assertEqual(self.ws.fetchTry('FIFO'), None)

    def testLifo(self):
        """A basic test of using a lifo mode variable"""
        self.ws.declare('LIFO', LIFO)
        values = range(10)

        self.log.info('Storing variables in LIFO mode')
        for i in values:
            self.ws.store('LIFO', i)

        self.log.info('Fetching values for LIFO variable')
        values.reverse()
        for i in values:
            t = self.ws.fetchTry('LIFO')
            self.assertEqual(t, i)

        self.assertEqual(self.ws.fetchTry('LIFO'), None)

    def testSingle(self):
        """A basic test of using a single mode variable"""
        self.ws.declare('SINGLE', SINGLE)
        values = range(10)

        self.log.info('Storing variables in SINGLE mode')
        for i in values:
            self.ws.store('SINGLE', i)

        self.log.info('Fetching values for SINGLE variable')
        self.assertEqual(self.ws.fetchTry('SINGLE'), values[-1])

        for i in values:
            self.assertEqual(self.ws.fetchTry('SINGLE'), None)

    def testCurrentWs(self):
        """Test the output of the currentWs method"""
        wsname = self.ws.currentWs()
        self.assertEqual(self.wsname, wsname)

    def testUseUseTrue(self):
        """Test setting useUse to True"""
        usews = NetWorkSpace(self.wsname, useUse=True, **nwsargs)
        val = 42
        self.ws.store('foo', val)
        self.assertEqual(usews.fetch('foo'), val)

    def testUseUseFalse(self):
        """Test that explicitly setting useUse=False works"""
        usews = NetWorkSpace(self.wsname + ' 2', useUse=False, **nwsargs)
        val = 7
        usews.store('bar', val)
        self.assertEqual(usews.findTry('bar'), val)
        self.assertEqual(usews.fetch('bar'), val)

    def testConstants(self):
        """Test the value of the mode constants"""
        self.assertEqual(FIFO, 'fifo')
        self.assertEqual(LIFO, 'lifo')
        self.assertEqual(MULTI, 'multi')
        self.assertEqual(SINGLE, 'single')

    def testDeleteWs(self):
        """Test deleting a workspace using deleteWs"""
        host = self.ws.server.serverHost
        port = self.ws.server.serverPort
        server = NwsServer(host, port)
        self.log.info('Deleting workspace %s' % repr(self.wsname))
        server.deleteWs(self.wsname)
        self.assertRaises(NwsOperationException, self.ws.store, 'fail', 0)

    def testListWss(self):
        """Test output from listWss"""
        wss = self.getWss(self.ws.server.listWss())
        self.assert_(wss.has_key('__default'))
        n = 0
        for k, v in wss.items():
            if v['mine']:
                n += 1
        self.assertEqual(n, 1)
        self.assert_(wss[self.wsname]['mine'])

    def testPersistent(self):
        """Test creation of persistent workspace"""
        wsname = 'persistent'
        t = NetWorkSpace(wsname, persistent=True, **nwsargs)
        try:
            wss = self.getWss(self.ws.server.listWss())
            self.assert_(wss.has_key(wsname))
            self.assert_(wss[wsname]['persistent'])
            t.server.close()

            # give the server a chance to notice that the connection was
            # dropped
            time.sleep(1)
            wss = self.getWss(self.ws.server.listWss())
            self.assert_(wss.has_key(wsname))
            self.assert_(wss[wsname]['persistent'])
        finally:
            self.ws.server.deleteWs(wsname)

        wss = self.getWss(self.ws.server.listWss())
        self.assert_(not wss.has_key(wsname))

    def testNonPersistent(self):
        """Test creation of non-persistent workspace"""
        wsname = 'non-persistent'
        t = NetWorkSpace(wsname, persistent=False, **nwsargs)
        wss = self.getWss(self.ws.server.listWss())
        self.assert_(wss.has_key(wsname))
        self.assert_(not wss[wsname]['persistent'])
        t.server.close()

        # give the server a chance to notice that the connection was
        # dropped
        time.sleep(1)
        wss = self.getWss(self.ws.server.listWss())
        self.assert_(not wss.has_key(wsname))

    def testMktempWs(self):
        """Test mktempWs"""
        self.log.info('calling mktempWs incorrectly...')
        self.assertRaises(NwsOperationException, self.ws.server.mktempWs, 'foo')

        self.log.info('calling mktempWs incorrectly again...')
        self.assertRaises(NwsOperationException, self.ws.server.mktempWs, 'foo')

        self.log.info('calling mktempWs a third time...')
        t = self.ws.server.mktempWs('foo%d')
        self.assert_(t.startswith('foo'))
        self.log.info('mktempWs returned: ' + t)
        # must explicitly delete, because nobody owns it
        self.ws.server.deleteWs(t)

    def testListVars(self):
        """Test listVars"""
        expected = {
            'foo': getdict(variable='foo', values=1, fetchers=0, finders=0, mode=SINGLE),
            'bar': getdict(variable='bar', values=2, fetchers=0, finders=0, mode=MULTI),
            'baz': getdict(variable='baz', values=1, fetchers=0, finders=0, mode=LIFO),
            'buz': getdict(variable='buz', values=0, fetchers=0, finders=0, mode=FIFO),
            'bif': getdict(variable='bif', values=0, fetchers=0, finders=0, mode=LIFO),
        }

        self.ws.declare('foo', SINGLE)
        self.ws.store('foo', 'hello, world')
        self.ws.store('foo', 'goodbye, world')
        self.ws.store('foo', 0)
        self.ws.store('foo', 222)
        self.ws.declare('bar', MULTI)
        self.ws.store('bar', 'spam')
        self.ws.store('bar', 'spam')
        self.ws.declare('baz', LIFO)
        self.ws.store('baz', 'eggs')
        self.ws.store('buz', 100)
        self.ws.fetchTry('buz')
        self.ws.declare('bif', LIFO)

        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)
        actual2 = self.getVarsDict(self.ws.listVars(format=DICT))
        self.assertEqual(actual2, expected)

        # now delete foo and bif
        self.ws.deleteVar('foo')
        self.ws.deleteVar('bif')
        del expected['foo']
        del expected['bif']
        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)
        actual2 = self.getVarsDict(self.ws.listVars(format=DICT))
        self.assertEqual(actual2, expected)

        # now delete baz
        self.ws.deleteVar('baz')
        del expected['baz']
        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)
        actual2 = self.getVarsDict(self.ws.listVars(format=DICT))
        self.assertEqual(actual2, expected)

        # now try the same thing again
        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)

        # now try to redeclare bar and buz
        self.assertRaises(NwsDeclarationFailedException, self.ws.declare, 'bar', FIFO)
        self.assertRaises(NwsDeclarationFailedException, self.ws.declare, 'buz', SINGLE)
        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)
        actual2 = self.getVarsDict(self.ws.listVars(format=DICT))
        self.assertEqual(actual2, expected)

        # now delete bar and buz
        self.ws.deleteVar('bar')
        self.ws.deleteVar('buz')
        del expected['bar']
        del expected['buz']
        actual = self.getVars(self.ws.listVars())
        self.assertEqual(actual, expected)
        actual2 = self.getVarsDict(self.ws.listVars(format=DICT))
        self.assertEqual(actual2, expected)
        self.assertEqual(len(actual), 0)
        self.assertEqual(len(actual2), 0)
        self.assertEqual(len(expected), 0)

    def testEmptyString(self):
        """Test that empty strings are handled correctly"""
        self.ws.store('Empty', '')
        self.assertEqual(self.ws.fetch('Empty'), '')
        self.assertEqual(self.ws.fetchTry('Empty'), None)

    def testDeclare(self):
        """Test variable declaration"""
        self.ws.declare('FIFO', FIFO)
        self.assertRaises(NwsDeclarationFailedException, self.ws.declare, 'FIFO', SINGLE)
        self.assertRaises(NwsDeclarationFailedException, self.ws.declare, 'FIFO', LIFO)
        self.assertRaises(NwsDeclarationFailedException, self.ws.declare, 'FIFO', MULTI)
        self.assertRaises(ValueError, self.ws.declare, 'FIFO', 'garbage')

    def testIfindTry(self):
        """Test ifindTry"""
        for i in range(10):
            self.ws.store('ifindTry', i)
        it = self.ws.ifindTry('ifindTry')
        self.assertEqual(list(it), range(10))
        self.assertRaises(StopIteration, it.next)
        self.assertEqual(list(it), [])
        it.restart()
        self.assertEqual(list(it), [])
        it.reset()
        self.assertEqual(list(it), range(10))
        self.assertEqual(list(it), [])
        it.restart()
        for i in range(10):
            self.ws.store('ifindTry', i + 10)
        self.assertEqual(list(it), range(10, 20))
        self.assertEqual(list(it), [])
        it.restart()
        self.assertEqual(list(it), [])
        it.reset()
        self.assertEqual(list(it), range(20))
        self.assertEqual(list(it), [])

    def testIfetchTry(self):
        """Test ifetchTry"""
        for i in range(10):
            self.ws.store('ifetchTry', i)
        it = self.ws.ifetchTry('ifetchTry')
        self.assertEqual(list(it), range(10))
        self.assertRaises(StopIteration, it.next)
        self.assertEqual(list(it), [])
        it.restart()
        self.assertEqual(list(it), [])
        self.assertEqual(list(it), [])
        it.reset()
        self.assertEqual(list(it), [])
        self.assertEqual(list(it), [])
        it.restart()
        for i in range(10):
            self.ws.store('ifetchTry', i + 10)
        self.assertEqual(list(it), range(10, 20))
        self.assertEqual(list(it), [])
        it.restart()
        self.assertEqual(list(it), [])
        self.assertEqual(list(it), [])
        it.reset()
        self.assertEqual(list(it), [])
        self.assertEqual(list(it), [])

    def testIfind(self):
        """Test ifind"""
        # XXX not a great test because it never blocks
        for i in range(10):
            self.ws.store('ifind', i)
        it = self.ws.ifind('ifind')
        self.assertEqual([it.next() for i in range(10)], range(10))
        for i in range(10):
            self.ws.store('ifind', i + 10)
        self.assertEqual([it.next() for i in range(10)], range(10, 20))
        it.reset()
        self.assertEqual([it.next() for i in range(20)], range(20))
        self.ws.deleteVar('ifind')
        self.assertRaises(StopIteration, it.next)
        self.assertEqual(list(it), [])

    def testIfetch(self):
        """Test ifetch"""
        # XXX not a great test because it never blocks
        for i in range(10):
            self.ws.store('ifetch', i)
        it = self.ws.ifetch('ifetch')
        self.assertEqual([it.next() for i in range(10)], range(10))
        for i in range(10):
            self.ws.store('ifetch', i + 10)
        self.assertEqual([it.next() for i in range(10)], range(10, 20))
        self.ws.deleteVar('ifetch')
        self.assertRaises(StopIteration, it.next)
        self.assertEqual(list(it), [])

    def testConcurrentIfetchTry(self):
        """Test multiple concurrent ifetchTry iterators"""
        for i in range(10):
            self.ws.store('ifetchTry', i)

        it_1 = self.ws.ifetchTry('ifetchTry')
        it_2 = self.ws.ifetchTry('ifetchTry')

        self.assertEqual(it_1.next(), 0)
        self.assertEqual(it_2.next(), 1)
        self.assertEqual(it_1.next(), 2)
        self.assertEqual(it_2.next(), 3)
        self.assertEqual(it_1.next(), 4)
        self.assertEqual(it_1.next(), 5)
        self.assertEqual(it_2.next(), 6)
        self.assertEqual(it_2.next(), 7)
        self.assertEqual(it_1.next(), 8)
        self.assertEqual(it_2.next(), 9)
        self.assertRaises(StopIteration, it_1.next)
        self.assertRaises(StopIteration, it_2.next)
        self.ws.deleteVar('ifetchTry')

    def testConcurrentIfindTry(self):
        """Test multiple concurrent ifindTry iterators"""
        for i in range(5):
            self.ws.store('conifindTry', i)

        it_1 = self.ws.ifindTry('conifindTry')
        it_2 = self.ws.ifindTry('conifindTry')

        self.assertEqual(it_1.next(), 0)
        self.assertEqual(it_1.next(), 1)
        self.assertEqual(it_2.next(), 0)
        self.assertEqual(it_2.next(), 1)
        self.assertEqual(it_1.next(), 2)
        self.assertEqual(it_2.next(), 2)
        self.assertEqual(it_1.next(), 3)
        self.assertEqual(it_1.next(), 4)
        self.assertEqual(it_2.next(), 3)
        self.assertEqual(it_2.next(), 4)
        self.assertRaises(StopIteration, it_1.next)
        self.assertRaises(StopIteration, it_2.next)
        self.ws.deleteVar('conifindTry')

    def testIfindTrySingle(self):
        """Test ifindTry with single mode variables"""
        self.ws.declare('ifindTrySingle', SINGLE)
        it = self.ws.ifindTry('ifindTrySingle')
        self.ws.store('ifindTrySingle', 100)
        self.assertEqual(it.next(), 100)
        self.ws.store('ifindTrySingle', 101)
        self.assertEqual(it.next(), 101)
        self.ws.store('ifindTrySingle', 102)
        self.ws.store('ifindTrySingle', 103)
        self.assertEqual(it.next(), 103)
        self.assertRaises(StopIteration, it.next)
        it.restart()
        self.ws.store('ifindTrySingle', 104)
        self.ws.store('ifindTrySingle', 105)
        self.assertEqual(it.next(), 105)
        self.assertRaises(StopIteration, it.next)
        it.reset()
        self.assertEqual(it.next(), 105)
        self.assertRaises(StopIteration, it.next)
        self.ws.deleteVar('ifindTrySingle')

    def testIfetchTrySingle(self):
        """Test ifetchTry with single mode variables"""
        self.ws.declare('ifetchTrySingle', SINGLE)
        it = self.ws.ifetchTry('ifetchTrySingle')
        self.ws.store('ifetchTrySingle', 100)
        self.assertEqual(it.next(), 100)
        self.ws.store('ifetchTrySingle', 101)
        self.assertEqual(it.next(), 101)
        self.ws.store('ifetchTrySingle', 102)
        self.ws.store('ifetchTrySingle', 103)
        self.assertEqual(it.next(), 103)
        self.assertRaises(StopIteration, it.next)
        it.restart()
        self.ws.store('ifetchTrySingle', 104)
        self.ws.store('ifetchTrySingle', 105)
        self.assertEqual(it.next(), 105)
        self.assertRaises(StopIteration, it.next)
        it.reset()
        self.assertRaises(StopIteration, it.next)
        self.ws.deleteVar('ifetchTrySingle')

    def testFetchTryFile(self):
        """Test concatenating string using fetchTryFile"""
        for i in string.ascii_lowercase:
            self.ws.store('fetchTryFile test', i)

        r = StringIO.StringIO()

        while self.ws.fetchTryFile('fetchTryFile test', r):
            pass

        self.assertEqual(r.getvalue(), string.ascii_lowercase)

    def testStoreFile(self):
        """Test splitting up a file using storeFile"""
        s = StringIO.StringIO(string.ascii_lowercase)
        while self.ws.storeFile('storeFile test', s, n=1):
            pass

        for i in string.ascii_lowercase:
            self.assertEqual(i, self.ws.fetchTry('storeFile test'))

    def testWriteTo(self):
        """Test writeTo method on object returned by ifindTry"""
        for i in string.ascii_lowercase:
            self.ws.store('writeTo', i)

        for i in range(3):
            r = StringIO.StringIO()
            it = self.ws.ifindTry('writeTo')
            while it.writeTo(r):
                pass
            self.assertEqual(r.getvalue(), string.ascii_lowercase)

        r = StringIO.StringIO()
        it = self.ws.ifetchTry('writeTo')
        while it.writeTo(r):
            pass
        self.assertEqual(r.getvalue(), string.ascii_lowercase)

        for i in range(3):
            r = StringIO.StringIO()
            it = self.ws.ifindTry('writeTo')
            while it.writeTo(r):
                pass
            self.assertEqual(r.getvalue(), "")

if __name__ == '__main__':
    import sys, os

    nwsargs = {}
    try: nwsargs['serverHost'] = os.environ['NWS_HOST']
    except: pass
    try: nwsargs['serverPort'] = int(os.environ['NWS_PORT'])
    except: pass
    verbose = os.environ.get('NWS_VERBOSE', False)

    basicConfig()
    log = getLogger('NWSTest')
    if verbose:
        log.setLevel(DEBUG)

    unittest.main()
