Mercurial > hg > GlobalNeighbors

comparison tests/test_write.py @ 22:e69cb496324e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
we have a data dump
author Jeff Hammel <k0scist@gmail.com>
date Sun, 25 Jun 2017 17:45:19 -0700
parents
children
comparison
equal deleted inserted replaced
21:22c384fe954d 22:e69cb496324e
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()