def test_process(self, mock_sink):
docs = [{'x': 1}, {'x': 2}, {'x': 3}]
with TestPipeline() as p:
_ = (p | "Create" >> beam.Create(docs)
| "Write" >> beam.ParDo(_WriteMongoFn(batch_size=2)))
p.run()
self.assertEqual(
2, mock_sink.return_value.__enter__.return_value.write.call_count)
def test_display_data_mask_password(self):
# Uri without password
data = _WriteMongoFn(uri='mongodb://test').display_data()
self.assertTrue('uri' in data)
self.assertTrue(data['uri'] == 'mongodb://test')
# Password is masked in the uri if present
data = _WriteMongoFn(
uri='mongodb+srv://user:[email protected]/testdb',
extra_params={
'user': '******',
'password': '******'
}).display_data()
self.assertTrue('uri' in data)
self.assertTrue(
data['uri'] == 'mongodb+srv://user:******@test.mongodb.net/testdb')
# Password is masked in the client spec if present
self.assertTrue('mongo_client_params' in data)
self.assertTrue(data['mongo_client_params'] ==
'{"user": "******", "password": "******"}')
def test_display_data(self):
data = _WriteMongoFn(batch_size=10).display_data()
self.assertEqual(10, data['batch_size'])
