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()