Skip to content

Instantly share code, notes, and snippets.

@serenamm
Last active September 9, 2019 18:41
Show Gist options
  • Save serenamm/a2d75abe5a043e6b2ce9c09c3774b8d0 to your computer and use it in GitHub Desktop.
Save serenamm/a2d75abe5a043e6b2ce9c09c3774b8d0 to your computer and use it in GitHub Desktop.
def test_create_new_table(mocker):
# Mock all our variables
mock_spark = mock.Mock()
mock_category_q = mock.Mock()
mock_created_table = mock.Mock()
mock_created_table_coalesced = mock.Mock()
# Calling spark.sql with create_table_query returns created_table - we need to mock it
mock_spark.sql.side_effect = [mock_created_table]
# Mock the output of calling .coalesce on created_table
mock_created_table.coalesce.return_value = mock_created_table_coalesced
# Mock the .write as well
mock_write = mock.Mock()
# Mock the output of calling .write on the coalesced created table
mock_created_table_coalesced.write = mock_write
test_paths = {
"product_table": {
"table": "products",
},
"similarity_table": {
"table": "product_similarity"
},
"created_table": {
"path": "path_to_table",
}
}
test_params = {
"num_items": 10,
}
# Call our function with our mocks
create_new_table(mock_spark, test_paths, test_params, mock_category_q)
# We only want spark.sql to have been called once, so assert that
assert 1 == mock_spark.sql.call_count
# Assert that we did in fact call created_table.coalesce(1)
mock_created_table.coalesce.assert_called_with(1)
# Assert that the table save path was passed in properly
mock_write.save.assert_called_with(test_paths["created_table"]["path"],
format="orc", mode="Overwrite")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment