import numpy as np
import pandas as pd

from spandex import TableFrame


def test_tableframe(loader):
    table = loader.database.tables.sample.hf_bg
    for cache in [False, True]:
        tf = TableFrame(table, index_name='gid', cache=cache)
        assert isinstance(tf.index, pd.Index)
        num_rows = len(tf)
        assert num_rows > 1
        assert set(tf.columns) == set(table.__table__.columns.keys())
        for column_name in tf.columns:
            if column_name != 'gid':
                if cache:
                    assert column_name not in tf._cached.keys()
                assert isinstance(tf[column_name], pd.Series)
                if cache:
                    assert column_name in tf._cached.keys()
                assert isinstance(getattr(tf, column_name), pd.Series)
        df = tf[['objectid']]
        assert isinstance(df, pd.DataFrame)
        assert len(df) == num_rows
        assert set(df.columns) == set(['objectid'])
        assert np.issubdtype(df.objectid.dtype, int)

@@ -36,39 +36,134 @@
  pd%0A
-%0Afrom spandex import TableFrame
+from pandas.util import testing as pdt%0Aimport pytest%0A%0Afrom spandex import TableFrame%0Afrom spandex.io import db_to_df, df_to_db
 %0A%0A%0Ad
@@ -213,17 +213,8 @@
 der.
-database.
 tabl
@@ -1093,8 +1093,1839 @@
 e, int)%0A
+%0A%0Adef test_sim_export(loader):%0A    # Try importing the UrbanSim simulation framework, otherwise skip test.%0A    sim = pytest.importorskip('urbansim.sim.simulation')%0A    from spandex.sim import column%0A%0A    # Register input parcels table.%0A    parcels = loader.tables.sample.heather_farms%0A    parcels_in = TableFrame(parcels, index_name='gid')%0A    sim.add_table('parcels_in', parcels_in, copy=False)%0A%0A    # Register output parcels table.%0A    @sim.table()%0A    def parcels_out(parcels_in):%0A        return pd.DataFrame(index=parcels_in.parcel_id)%0A%0A    # Specify default table for output columns as decorator.%0A    def out(*args, **kwargs):%0A        return column(table_name='parcels_out', groupby=parcels_in.parcel_id,%0A                      *args, **kwargs)%0A%0A    # Specify some output columns.%0A    @out(astype=str, agg='first')%0A    def apn(apn='parcels_in.puid'):%0A        return apn%0A%0A    @out()%0A    def county_id():%0A        return 13%0A%0A    @out(astype=float, agg='median')%0A    def area(acr='parcels_in.parcel_acr'):%0A        return 4047. * acr%0A%0A    # Register model to export output table to database.%0A    @sim.model()%0A    def export(parcels_out):%0A        schema = loader.tables.sample%0A        df_to_db(parcels_out.to_frame(), 'parcels_out', schema=schema)%0A%0A    # Inspect output table.%0A    column_names = %5B'apn', 'county_id', 'area'%5D%0A    parcels_out_df1 = sim.get_table('parcels_out').to_frame()%0A    assert set(parcels_out_df1.columns) == set(column_names)%0A    assert parcels_out_df1.county_id.unique() == %5B13%5D%0A%0A    # Export table to database and import back to compare.%0A    sim.run(%5B'export'%5D)%0A    parcels_out_table = loader.tables.sample.parcels_out%0A    parcels_out_df2 = db_to_df(parcels_out_table, index_name='parcel_id')%0A    pdt.assert_frame_equal(parcels_out_df1%5Bcolumn_names%5D,%0A                           parcels_out_df2%5Bcolumn_names%5D)%0A

Add test for UrbanSim simulation framework