added setup_tests.py

fmcclean · Sep 9, 2020 · 96317de · 96317de
1 parent a3d4760
commit 96317de
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 88 deletions.
diff --git a/citycatpg/run.py b/citycatpg/run.py
@@ -191,13 +191,13 @@ def get_rainfall_polygons(self, con):
         else:
             return gpd.GeoDataFrame.from_postgis(
                 sql.SQL(
-                """
-                SELECT {rain_geom_table}.gid, {rain_geom_table}.geom FROM {rain_geom_table}, {domain_table} 
-                WHERE ST_Intersects({rain_geom_table}.geom, {domain_table}.geom)
-                """).format(
-                    rain_geom_table=sql.Identifier(self.rain_geom_table),
-                    domain_table=sql.Identifier(self.domain_table)
-                ).as_string(con),
+                    """
+                    SELECT {rain_geom_table}.gid, {rain_geom_table}.geom FROM {rain_geom_table}, {domain_table} 
+                    WHERE ST_Intersects({rain_geom_table}.geom, {domain_table}.geom)
+                    """).format(
+                        rain_geom_table=sql.Identifier(self.rain_geom_table),
+                        domain_table=sql.Identifier(self.domain_table)
+                    ).as_string(con),
                 con)
 
     def execute(self, run_path, out_path):

diff --git a/tests/setup_tests.py b/tests/setup_tests.py
@@ -0,0 +1,64 @@
+from shapely.geometry import Polygon
+import rasterio as rio
+from rasterio.transform import Affine
+import numpy as np
+import psycopg2
+from psycopg2 import sql
+from citycatpg import Run
+
+con = psycopg2.connect(database='postgres', user='postgres', password='password', host='localhost')
+
+dem_file = rio.MemoryFile()
+x_min, y_max, res, height, width = 100, 500, 5, 100, 200
+x_max, y_min = x_min + width * res, y_max - height * res
+array = np.round(np.random.random((height, width)), 3)
+
+with rio.open(
+        dem_file,
+        'w',
+        driver='GTiff',
+        height=height,
+        width=width,
+        count=1,
+        dtype=array.dtype,
+        transform=Affine.translation(x_min, y_max) * Affine.scale(res, -res),
+        nodata=-9999
+) as dst:
+    dst.write(array, 1)
+
+dem_file.seek(0)
+
+r = Run(100, 3035, 90, rain_table='rain')
+
+with con:
+    with con.cursor() as cursor:
+        cursor.execute("CREATE EXTENSION IF NOT EXISTS postgis_raster")
+        cursor.execute("SET postgis.gdal_enabled_drivers TO 'GTiff'")
+
+        cursor.execute(
+            sql.SQL("""
+                        DROP TABLE IF EXISTS {domain_table};
+                        CREATE TABLE {domain_table} (gid serial PRIMARY KEY, geom geometry);
+                        INSERT INTO {domain_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
+
+                        DROP TABLE IF EXISTS {rain_geom_table};
+                        CREATE TABLE {rain_geom_table} (gid serial PRIMARY KEY, geom geometry);
+                        INSERT INTO {rain_geom_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
+
+                        DROP TABLE IF EXISTS {rain_table};
+                        CREATE TABLE {rain_table} (gid integer, time timestamp, value numeric);
+                        INSERT INTO {rain_table} (gid, time, value) VALUES (1, '2000-01-01', 10),(1, '2000-01-02', 20);
+
+                        DROP TABLE IF EXISTS {dem_table};
+                        CREATE TABLE {dem_table} (rast raster);
+                        INSERT INTO {dem_table}(rast) VALUES (ST_FromGDALRaster(%(rast)s));
+                    """).format(
+                domain_table=sql.Identifier(r.domain_table),
+                dem_table=sql.Identifier(r.dem_table),
+                rain_geom_table=sql.Identifier(r.rain_geom_table),
+                rain_table=sql.Identifier(r.rain_table),
+            ),
+            dict(
+                geom=str(Polygon(
+                    [[x_min, y_min], [x_min, y_max], [x_max / 2, y_max / 2], [x_max, y_min], [x_min, y_min]])),
+                rast=psycopg2.Binary(dem_file.read())))
diff --git a/tests/test_run.py b/tests/test_run.py
@@ -1,81 +1,17 @@
 from citycatpg import Run, fetch
 from unittest import TestCase
-from shapely.geometry import Polygon
-import rasterio as rio
-from rasterio.transform import Affine
-import numpy as np
-import psycopg2
-from psycopg2 import sql
+from .setup_tests import con
 import datetime
 
-con = psycopg2.connect(database='postgres', user='postgres', password='password', host='localhost')
-
 
 class TestRun(TestCase):
 
-    @classmethod
-    def setUpClass(cls):
-        """Insert sample data"""
-        dem_file = rio.MemoryFile()
-        x_min, y_max, res, height, width = 100, 500, 5, 100, 200
-        x_max, y_min = x_min + width * res, y_max - height * res
-        array = np.round(np.random.random((height, width)), 3)
-
-        with rio.open(
-                dem_file,
-                'w',
-                driver='GTiff',
-                height=height,
-                width=width,
-                count=1,
-                dtype=array.dtype,
-                transform=Affine.translation(x_min, y_max) * Affine.scale(res, -res),
-                nodata=-9999
-        ) as dst:
-            dst.write(array, 1)
-
-        dem_file.seek(0)
-
-        r = Run(100, 3035, 90, rain_table='rain')
-
-        with con:
-            with con.cursor() as cursor:
-                cursor.execute("CREATE EXTENSION IF NOT EXISTS postgis_raster")
-                cursor.execute("SET postgis.gdal_enabled_drivers TO 'GTiff'")
-
-                cursor.execute(
-                    sql.SQL("""
-                        DROP TABLE IF EXISTS {domain_table};
-                        CREATE TABLE {domain_table} (gid serial PRIMARY KEY, geom geometry);
-                        INSERT INTO {domain_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
-
-                        DROP TABLE IF EXISTS {rain_geom_table};
-                        CREATE TABLE {rain_geom_table} (gid serial PRIMARY KEY, geom geometry);
-                        INSERT INTO {rain_geom_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
-                        
-                        DROP TABLE IF EXISTS {rain_table};
-                        CREATE TABLE {rain_table} (gid integer, time timestamp, value numeric);
-                        INSERT INTO {rain_table} (gid, time, value) VALUES (1, '2000-01-01', 10),(1, '2000-01-02', 20);
-
-                        DROP TABLE IF EXISTS {dem_table};
-                        CREATE TABLE {dem_table} (rast raster);
-                        INSERT INTO {dem_table}(rast) VALUES (ST_FromGDALRaster(%(rast)s));
-                    """).format(
-                        domain_table=sql.Identifier(r.domain_table),
-                        dem_table=sql.Identifier(r.dem_table),
-                        rain_geom_table=sql.Identifier(r.rain_geom_table),
-                        rain_table=sql.Identifier(r.rain_table),
-                    ),
-                    dict(
-                        geom=str(Polygon(
-                            [[x_min, y_min], [x_min, y_max], [x_max / 2, y_max / 2], [x_max, y_min], [x_min, y_min]])),
-                        rast=psycopg2.Binary(dem_file.read())))
-
     def test_add_run(self):
         with con:
             with con.cursor() as cur:
                 cur.execute('DROP TABLE IF EXISTS runs')
-        run = Run(run_duration=500, srid=3035, resolution=90, rain_total=100, rain_duration=200, run_name='test')
+        run = Run(run_duration=120, srid=3035, resolution=90, rain_total=100, rain_duration=120, run_name='test',
+                  output_frequency=60)
         run.add(con)
         return run
 
@@ -84,20 +20,25 @@ def test_fetch(self):
             with con.cursor() as cur:
                 cur.execute('DROP TABLE IF EXISTS runs')
 
-        r = Run(run_duration=500, srid=3035, resolution=90)
-        r.add(con)
+        run = Run(run_duration=500, srid=3035, resolution=90)
+        run.add(con)
 
-        fetch(con, run_id=r.run_id)
+        fetch(con, run_id=run.run_id)
 
     def test_generate_rainfall(self):
-        r = Run(100, 3035, 90, rain_total=100, rain_duration=500)
-        r.get_model(con)
-        r.model.write('tests/test_model_generated')
+        run = Run(100, 3035, 90, rain_total=100, rain_duration=500)
+        run.get_model(con)
+        run.model.write('tests/test_model_generated')
 
     def test_get_model(self):
-        r = Run(100, 3035, 90, rain_table='rain',
-                rain_start=datetime.datetime(2000, 1, 1),
-                rain_end=datetime.datetime(2000, 1, 2))
+        run = Run(100, 3035, 90, rain_table='rain',
+                  rain_start=datetime.datetime(2000, 1, 1), rain_end=datetime.datetime(2000, 1, 2))
+
+        run.get_model(con)
+        run.model.write('tests/test_model')
 
-        r.get_model(con)
-        r.model.write('tests/test_model')
+    def test_execute(self):
+        run = Run(run_duration=120, srid=3035, resolution=90, rain_total=100, rain_duration=100, run_name='test',
+                  output_frequency=60)
+        run.get_model(con)
+        run.execute('tests/test_model_execute', 'tests/test_model_execute')
diff --git a/tests/test_server.py b/tests/test_server.py
@@ -1,8 +1,9 @@
 from citycatpg import server
+from citycatpg.run import Run
 from unittest import TestCase
 import pika
-from .test_run import TestRun, con
 import os
+from .setup_tests import con
 
 
 class TestServer(TestCase):
@@ -13,8 +14,9 @@ def test_run_server(self):
         channel = connection.channel()
         channel.queue_declare(queue=queue)
 
-        test_run = TestRun()
-        run = test_run.test_add_run()
+        run = Run(run_duration=120, srid=3035, resolution=90, rain_total=100, rain_duration=120, run_name='test',
+                  output_frequency=60)
+        run.add(con)
 
         channel.basic_publish(exchange='',
                               routing_key=queue,