def test_wait_no_instance(self, mock_hook): mock_hook.return_value.get_instance.return_value = None op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id", gcp_conn_id=GCP_CONN_ID) self.assertFalse(op.poke(None)) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
def test_wait_ready(self, mock_hook): mock_hook.return_value.get_instance.return_value = mock.Mock(Instance) mock_hook.return_value.get_cluster_states_for_table.return_value = { "cl-id": ClusterState(4) } op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id", gcp_conn_id=GCP_CONN_ID) self.assertTrue(op.poke(None)) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
def test_wait_no_table(self, mock_hook): mock_hook.return_value.get_instance.return_value = mock.Mock(Instance) mock_hook.return_value.get_cluster_states_for_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.NotFound( "Table not found.")) op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id", gcp_conn_id=GCP_CONN_ID) self.assertFalse(op.poke(None)) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
def test_wait_no_instance(self, mock_hook): mock_hook.return_value.get_instance.return_value = None op = BigtableTableReplicationCompletedSensor( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id", gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN, ) assert not op.poke(None) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN, )
def test_wait_not_ready(self, mock_hook): mock_hook.return_value.get_instance.return_value = mock.Mock(Instance) mock_hook.return_value.get_cluster_states_for_table.return_value = { "cl-id": ClusterState(0) } op = BigtableTableReplicationCompletedSensor( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id", gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN, ) assert not op.poke(None) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN, )