import unittest
from pysus.demography import geobase
import geopandas as gpd

from time import time

class Geobase(unittest.TestCase):
    def test_levels(self):
        self.assertIn('Municipality', geobase.LEVELS)

    def test_get_map(self):
        GB = geobase.GeoBase(level='Municipality')
        gdf = GB.map(4108304)
        self.assertIsInstance(gdf, gpd.GeoDataFrame)
        t0 = time()
        gdf = GB.map() # from second time it should skip download
        tf = time()
        self.assertLess(tf-t0, 1)
        self.assertEqual(int(gdf['code_muni'].values[0]), 4108304)
        self.assertEqual(len(gdf), 1)

    def test_get_map_2019(self):
        GB = geobase.GeoBase(level='Municipality', reset=True)
        gdf = GB.map(4220000, year=2019)
        self.assertIsInstance(gdf, gpd.GeoDataFrame)
        t0 = time()
        gdf = GB.map()  # from second time it should skip download
        tf = time()
        self.assertLess(tf - t0, 1)
        self.assertEqual(int(gdf['code_muni'].values[0]), 4220000)
        self.assertEqual(len(gdf), 1)

    def test_demographics(self):
        GB = geobase.GeoBase(level='State')
        gdf = GB.map('all')
        self.assertEqual(len(gdf), 27)
        GB.demographics()
        self.assertIn('population', GB.mapdf)



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