class TestGCSUploadSessionCompleteSensor(TestCase): def setUp(self): args = { 'owner': 'airflow', 'start_date': DEFAULT_DATE, } dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args) dag.schedule_interval = '@once' self.dag = dag self.sensor = GCSUploadSessionCompleteSensor( task_id='sensor_1', bucket='test-bucket', prefix='test-prefix/path', inactivity_period=12, poke_interval=10, min_objects=1, allow_delete=False, google_cloud_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO, impersonation_chain=TEST_IMPERSONATION_CHAIN, dag=self.dag, ) self.last_mocked_date = datetime(2019, 4, 24, 0, 0, 0) @mock.patch("airflow.providers.google.cloud.sensors.gcs.GCSHook") def test_get_gcs_hook(self, mock_hook): self.sensor._get_gcs_hook() mock_hook.assert_called_once_with( gcp_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) assert mock_hook.return_value == self.sensor.hook @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_files_deleted_between_pokes_throw_error(self): self.sensor.is_bucket_updated({'a', 'b'}) with pytest.raises(AirflowException): self.sensor.is_bucket_updated({'a'}) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_files_deleted_between_pokes_allow_delete(self): self.sensor = GCSUploadSessionCompleteSensor( task_id='sensor_2', bucket='test-bucket', prefix='test-prefix/path', inactivity_period=12, poke_interval=10, min_objects=1, allow_delete=True, dag=self.dag, ) self.sensor.is_bucket_updated({'a', 'b'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a'}) assert len(self.sensor.previous_objects) == 1 assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a', 'c'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a', 'd'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a', 'd'}) assert self.sensor.inactivity_seconds == 10 assert self.sensor.is_bucket_updated({'a', 'd'}) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_incoming_data(self): self.sensor.is_bucket_updated({'a'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a', 'b'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a', 'b', 'c'}) assert self.sensor.inactivity_seconds == 0 @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_no_new_data(self): self.sensor.is_bucket_updated({'a'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a'}) assert self.sensor.inactivity_seconds == 10 @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_no_new_data_success_criteria(self): self.sensor.is_bucket_updated({'a'}) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated({'a'}) assert self.sensor.inactivity_seconds == 10 assert self.sensor.is_bucket_updated({'a'}) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_not_enough_objects(self): self.sensor.is_bucket_updated(set()) assert self.sensor.inactivity_seconds == 0 self.sensor.is_bucket_updated(set()) assert self.sensor.inactivity_seconds == 10 assert not self.sensor.is_bucket_updated(set())
class TestGCSUploadSessionCompleteSensor(TestCase): def setUp(self): args = { 'owner': 'airflow', 'start_date': DEFAULT_DATE, } dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args) dag.schedule_interval = '@once' self.dag = dag self.sensor = GCSUploadSessionCompleteSensor( task_id='sensor_1', bucket='test-bucket', prefix='test-prefix/path', inactivity_period=12, poke_interval=10, min_objects=1, allow_delete=False, dag=self.dag ) self.last_mocked_date = datetime(2019, 4, 24, 0, 0, 0) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_files_deleted_between_pokes_throw_error(self): self.sensor.is_bucket_updated({'a', 'b'}) with self.assertRaises(AirflowException): self.sensor.is_bucket_updated({'a'}) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_files_deleted_between_pokes_allow_delete(self): self.sensor = GCSUploadSessionCompleteSensor( task_id='sensor_2', bucket='test-bucket', prefix='test-prefix/path', inactivity_period=12, poke_interval=10, min_objects=1, allow_delete=True, dag=self.dag ) self.sensor.is_bucket_updated({'a', 'b'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a'}) self.assertEqual(len(self.sensor.previous_objects), 1) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a', 'c'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a', 'd'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a', 'd'}) self.assertEqual(self.sensor.inactivity_seconds, 10) self.assertTrue(self.sensor.is_bucket_updated({'a', 'd'})) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_incoming_data(self): self.sensor.is_bucket_updated({'a'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a', 'b'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a', 'b', 'c'}) self.assertEqual(self.sensor.inactivity_seconds, 0) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_no_new_data(self): self.sensor.is_bucket_updated({'a'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a'}) self.assertEqual(self.sensor.inactivity_seconds, 10) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_no_new_data_success_criteria(self): self.sensor.is_bucket_updated({'a'}) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated({'a'}) self.assertEqual(self.sensor.inactivity_seconds, 10) self.assertTrue(self.sensor.is_bucket_updated({'a'})) @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time) def test_not_enough_objects(self): self.sensor.is_bucket_updated(set()) self.assertEqual(self.sensor.inactivity_seconds, 0) self.sensor.is_bucket_updated(set()) self.assertEqual(self.sensor.inactivity_seconds, 10) self.assertFalse(self.sensor.is_bucket_updated(set()))