|
22
|
1 #!/usr/bin/env python
|
|
|
2
|
|
|
3 """
|
|
|
4 test writing + reading distances
|
|
|
5 """
|
|
|
6
|
|
|
7 import os
|
|
|
8 import shutil
|
|
|
9 import tempfile
|
|
|
10 import unittest
|
|
|
11 from common import datafile
|
|
|
12 from globalneighbors.distance import calculate_neighbors
|
|
|
13 from globalneighbors.distance import write_neighbors
|
|
|
14 from globalneighbors.locations import locations
|
|
|
15 from globalneighbors.neighbors import read_neighbors_file
|
|
|
16 from globalneighbors.read import read_cities
|
|
|
17
|
|
|
18
|
|
|
19 class TestDistanceReadWrite(unittest.TestCase):
|
|
|
20
|
|
|
21 def test_10000(self):
|
|
|
22 """test 10000 cities"""
|
|
|
23
|
|
|
24 # read locations
|
|
|
25 citiesfile = datafile('10000cities.tsv')
|
|
|
26 assert os.path.exists(citiesfile)
|
|
|
27 with open(citiesfile) as f:
|
|
|
28 city_locations = locations(read_cities(f))
|
|
|
29
|
|
|
30 # calculate neighbors
|
|
|
31 neighbors = calculate_neighbors(city_locations,
|
|
|
32 k=50,
|
|
|
33 lat_tol=2.,
|
|
|
34 lon_tol=2.)
|
|
|
35
|
|
|
36 # make a staging area
|
|
|
37 tmpdir = tempfile.mkdtemp()
|
|
|
38 try:
|
|
|
39 # write the neighbors
|
|
|
40 outfile = os.path.join(tmpdir, 'neighbors.dat')
|
|
|
41 assert not os.path.exists(outfile)
|
|
|
42 with open(outfile, 'w') as f:
|
|
|
43 write_neighbors(f, neighbors)
|
|
|
44 assert os.path.exists(outfile)
|
|
|
45
|
|
|
46 # read the neighbors
|
|
|
47 with open(outfile) as f:
|
|
|
48 new_neighbors = read_neighbors_file(f)
|
|
|
49 finally:
|
|
|
50 shutil.rmtree(tmpdir, ignore_errors=True)
|
|
|
51
|
|
|
52 # they should be equal
|
|
|
53 assert len(neighbors) == len(new_neighbors)
|
|
|
54 assert sorted(neighbors.keys()) == sorted(new_neighbors.keys())
|
|
|
55 for key in neighbors.keys():
|
|
|
56 valueA = neighbors[key]
|
|
|
57 valueB = new_neighbors[key]
|
|
|
58 assert valueA == valueB
|
|
|
59
|
|
|
60
|
|
|
61 if __name__ == '__main__':
|
|
|
62 unittest.main()
|
