Last active
September 9, 2019 18:41
-
-
Save serenamm/a2d75abe5a043e6b2ce9c09c3774b8d0 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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