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'])