def test_poke(self, mock_hook): sensor = CassandraRecordSensor(task_id='test_task', cassandra_conn_id='cassandra_default', table='t', keys={'foo': 'bar'}) sensor.poke(None) mock_hook.return_value.record_exists.assert_called_once_with( 't', {'foo': 'bar'})
def setUp(self): args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} self.dag = DAG('test_dag_id', default_args=args) self.sensor = CassandraRecordSensor( task_id='test_task', cassandra_conn_id='cassandra_default', dag=self.dag, table='t', keys={'foo': 'bar'})
def test_poke_should_not_fail_with_empty_keys(self, mock_hook): sensor = CassandraRecordSensor( task_id='test_task', cassandra_conn_id=TEST_CASSANDRA_CONN_ID, table=TEST_CASSANDRA_TABLE, keys=None, ) exists = sensor.poke(dict()) self.assertTrue(exists) mock_hook.return_value.record_exists.assert_called_once_with( TEST_CASSANDRA_TABLE, None) mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
def test_poke(self, mock_hook): sensor = CassandraRecordSensor( task_id='test_task', cassandra_conn_id=TEST_CASSANDRA_CONN_ID, table=TEST_CASSANDRA_TABLE, keys=TEST_CASSANDRA_KEY, ) exists = sensor.poke(dict()) self.assertTrue(exists) mock_hook.return_value.record_exists.assert_called_once_with( TEST_CASSANDRA_TABLE, TEST_CASSANDRA_KEY) mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
class TestCassandraRecordSensor(unittest.TestCase): def setUp(self): args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} self.dag = DAG('test_dag_id', default_args=args) self.sensor = CassandraRecordSensor( task_id='test_task', cassandra_conn_id='cassandra_default', dag=self.dag, table='t', keys={'foo': 'bar'}) @patch("airflow.contrib.hooks.cassandra_hook.CassandraHook.record_exists") def test_poke(self, mock_record_exists): self.sensor.poke(None) mock_record_exists.assert_called_once_with('t', {'foo': 'bar'})
def test_poke_should_return_false_for_non_existing_table(self, mock_hook): mock_hook.return_value.record_exists.return_value = False sensor = CassandraRecordSensor( task_id='test_task', cassandra_conn_id=TEST_CASSANDRA_CONN_ID, table=TEST_CASSANDRA_TABLE, keys=TEST_CASSANDRA_KEY, ) exists = sensor.poke(dict()) assert not exists mock_hook.return_value.record_exists.assert_called_once_with( TEST_CASSANDRA_TABLE, TEST_CASSANDRA_KEY) mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
# KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """ Example Airflow DAG to check if a Cassandra Table and a Records exists or not using `CassandraTableSensor` and `CassandraRecordSensor`. """ from datetime import datetime from airflow.models import DAG from airflow.providers.apache.cassandra.sensors.record import CassandraRecordSensor from airflow.providers.apache.cassandra.sensors.table import CassandraTableSensor # [START howto_operator_cassandra_sensors] with DAG( dag_id='example_cassandra_operator', schedule_interval=None, start_date=datetime(2021, 1, 1), default_args={'table': 'keyspace_name.table_name'}, catchup=False, tags=['example'], ) as dag: table_sensor = CassandraTableSensor(task_id="cassandra_table_sensor") record_sensor = CassandraRecordSensor(task_id="cassandra_record_sensor", keys={ "p1": "v1", "p2": "v2" }) # [END howto_operator_cassandra_sensors]
args = { 'owner': 'Airflow', } with DAG( dag_id='example_cassandra_operator', default_args=args, schedule_interval=None, start_date=days_ago(2), tags=['example'], ) as dag: # [START howto_operator_cassandra_table_sensor] table_sensor = CassandraTableSensor( task_id="cassandra_table_sensor", cassandra_conn_id="cassandra_default", table="keyspace_name.table_name", ) # [END howto_operator_cassandra_table_sensor] # [START howto_operator_cassandra_record_sensor] record_sensor = CassandraRecordSensor( task_id="cassandra_record_sensor", cassandra_conn_id="cassandra_default", table="keyspace_name.table_name", keys={ "p1": "v1", "p2": "v2" }, ) # [END howto_operator_cassandra_record_sensor]