Mercurial > hg > GlobalNeighbors
diff tests/test_write.py @ 22:e69cb496324e
we have a data dump
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Sun, 25 Jun 2017 17:45:19 -0700 |
parents | |
children |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/tests/test_write.py Sun Jun 25 17:45:19 2017 -0700 @@ -0,0 +1,62 @@ +#!/usr/bin/env python + +""" +test writing + reading distances +""" + +import os +import shutil +import tempfile +import unittest +from common import datafile +from globalneighbors.distance import calculate_neighbors +from globalneighbors.distance import write_neighbors +from globalneighbors.locations import locations +from globalneighbors.neighbors import read_neighbors_file +from globalneighbors.read import read_cities + + +class TestDistanceReadWrite(unittest.TestCase): + + def test_10000(self): + """test 10000 cities""" + + # read locations + citiesfile = datafile('10000cities.tsv') + assert os.path.exists(citiesfile) + with open(citiesfile) as f: + city_locations = locations(read_cities(f)) + + # calculate neighbors + neighbors = calculate_neighbors(city_locations, + k=50, + lat_tol=2., + lon_tol=2.) + + # make a staging area + tmpdir = tempfile.mkdtemp() + try: + # write the neighbors + outfile = os.path.join(tmpdir, 'neighbors.dat') + assert not os.path.exists(outfile) + with open(outfile, 'w') as f: + write_neighbors(f, neighbors) + assert os.path.exists(outfile) + + # read the neighbors + with open(outfile) as f: + new_neighbors = read_neighbors_file(f) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + # they should be equal + assert len(neighbors) == len(new_neighbors) + assert sorted(neighbors.keys()) == sorted(new_neighbors.keys()) + for key in neighbors.keys(): + valueA = neighbors[key] + valueB = new_neighbors[key] + assert valueA == valueB + + +if __name__ == '__main__': + unittest.main()