1

I have a function to read data from a MySQL database:

def example(mysql_engine) -> DataFrame:
    query = """SELECT col_1 FROM xxx.xxx"""
    df = pandas.read_sql(query, mysql_engine)

    return df

mysql_engine is returned by another function, now I want to write unit tests to validate the column name col_1 and datatype in df, is there an example I can follow? Do I need to set up a real database engine in unit test?

wawawa
  • 2,835
  • 6
  • 44
  • 105
  • Assert that the df columns are ['id', 'blah', 'foo', 'bar'] or `df.dtypes` matches what you expect? – ifly6 Jul 28 '21 at 13:40
  • both, so one test case is to validate `col_1` is in `df`, another test case is to validate the dtypes of `col_1` is `int64`(for instance). – wawawa Jul 28 '21 at 13:45

1 Answers1

0

You can use assert col_1 in df.columns for the unit tests.

For the Database engine, you can use a mock object. You can check this link from the standard library.

S. Ali Mirferdos
  • 192
  • 1
  • 2
  • 12
  • but how should I get the `df` in the unit test? I'm thinking if I define a `dummy_df = {'col_1': ['123456', '789456']}` - assuming this is the dataframe that is returned from `read_sql`, then use `mock_read_sql.return_value = dummy_df` then `self.assertTrue('col_1' in dummy_df.columns)`, but something doesn't feel right about this approach... I don't know how to get `df` so that I can use `assertTrue('col_1' in df.columns)` instead of `assertTrue('col_1' in dummy_df.columns)` – wawawa Jul 28 '21 at 14:35
  • @Cecilia You actually want to unit test the function which the `df` is the output of it. You can pass a mock engine where mock.return_value equals the result you want from the DB. Then you would call your function and assert the output `df` – S. Ali Mirferdos Jul 28 '21 at 14:45
  • ` Then you would call your function and assert the output df` - if I call the function, does that mean I have to query against the real database? Don't think I want to do that, if I don't query against the real db with the query, unit test will complain... Is there an example? – wawawa Jul 28 '21 at 15:22