"""
Unit tests for `dbgroup` module.
"""
__author__ = "Dan Gunter"
__copyright__ = "Copyright 2014, The Materials Project"
__email__ = "dkgunter@lbl.gov"
__date__ = "2014-04-29"

import json
import mongomock
import os
import tempfile
import unittest
from matgendb.dbgroup import ConfigGroup
from matgendb import dbconfig

_opj = os.path.join

mockdb = mongomock.MongoClient()
doc = {"hello": "world"}
mockdb.testdb.data.insert(doc)
# add some nested collections
mockcoll = [
    'data.a1',
    'data.a1.b1',
    'data.a1.b2',
    'data.a2'
]
[mockdb.testdb[c].insert(doc) for c in mockcoll]

class MockQueryEngine:
    def __init__(self, **kwargs):
        self.kw = kwargs

    def __eq__(self, other):
        return other.kw == self.kw

    @property
    def db(self):
        return mockdb.testdb

    @property
    def collection(self):
        return mockdb.testdb[self.kw['collection']]

    @property
    def db(self):
        class DB:
            _c = mockcoll
            def collection_names(self, x=None):
                return self._c
        return DB()

class Cfg:
    def __init__(self, v):
        self.settings = {'collection': v}
        self.collection = v
    def copy(self):
        return Cfg(self.collection)

class ConfigGroupTestCase(unittest.TestCase):
    def setUp(self):
        self.g = ConfigGroup(qe_class=MockQueryEngine)
        self.configs = [Cfg("qe{:d}".format(i)) for i in range(5)]

    def test_add(self):
        """ConfigGroup add and lookup
        """
        keys = ["foo", "bar", "foo.a", "foo.b"]
        expect = {}
        for i, k in enumerate(keys):
            self.g.add(k, self.configs[i])
            expect[k] = MockQueryEngine(**self.configs[i].settings)
        self.assertEqual(self.g["foo"], expect["foo"])
        self.assertEqual(self.g["bar"], expect["bar"])
        self.assertEqual(self.g["bar*"], {"bar": expect["bar"]})
        self.assertEqual(self.g["foo.a"], expect["foo.a"])
        self.assertEqual(self.g["foo.*"],
                         {"foo.a": expect["foo.a"], "foo.b": expect["foo.b"]})

    def test_add_path(self):
        """Add set of query engines from a path.
        """
        # directory of pretend configs
        d = tempfile.mkdtemp()
        try:
            # fill with some configs
            c = {}
            for root in ("foo", "bar"):
                for sub in ("a", "b.1", "b.2"):
                    config = {
                        dbconfig.DB_KEY: root,
                        dbconfig.COLL_KEY: sub
                    }
                    filename = "mg_core_{}_{}.json".format(root, sub)
                    with open(_opj(d, filename), "w") as fp:
                        json.dump(config, fp)
                    c["{}.{}".format(root, sub)] = config
            # read them
            self.g.add_path(d)
            # check all were added
            self.assertEqual(sorted(self.g.keys()), sorted(c.keys()))
            # check one
            qe1 = self.g["foo.a"].kw
            c1 = c["foo.a"]
            self.assertTrue(dict_subset(c1, qe1))
            # check with prefix
            self.g.set_prefix("foo.b")
            self.assertTrue(dict_subset(c["foo.b.1"], self.g["1"].kw))
            self.assertRaises(KeyError, self.g.__getitem__, "bla")
            # check list with prefix
            gkeys = sorted(self.g["*"].keys())
            self.assertEqual(gkeys, ["foo.b.1", "foo.b.2"])
            # check list w/o prefix
            self.g.set_prefix()
            gkeys = sorted(self.g["bar.b.*"].keys())
            self.assertEqual(gkeys, ["bar.b.1", "bar.b.2"])
        finally:
            # rm -r $d
            for f in os.listdir(d):
                os.unlink(os.path.join(d, f))
            os.rmdir(d)

    def test_uncache(self):
        """Remove cached query engine(s) from ConfigGroup.
        """
        keys = ("foo.a", "foo", "bar")
        for i in range(len(keys)):
            self.g.add(keys[i], self.configs[i])
        # force instantiation/caching
        for i in range(len(keys)):
            self.g[keys[i]]
        left_behind = self.g[keys[2]]
        # remove all foo from cache
        self.g.uncache("foo*")
        # check that they are not cached
        for i in range(2):
            self.assertRaises(KeyError, self.g._cached.__getitem__, keys[i])
        # check that un-removed remain
        self.assertEqual(self.g[keys[2]], left_behind)

    def test_expand(self):
        """Add multiple collections at once with 'expand'.
        """
        self.g.add("foo", Cfg("data"), expand=True)
        # check that data.* got added as foo.*
        keys = set(self.g.keys())
        expect = set(["foo"] + [f.replace("data", "foo") for f in mockcoll])
        self.assertEqual(expect, keys)

def dict_subset(a, b):
    for k in a.keys():
        if k not in b or b[k] != a[k]:
            return False
    return True

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