def test_wait_no_instance(self, mock_hook):
        mock_hook.return_value.get_instance.return_value = None

        op = BigtableTableWaitForReplicationSensor(project_id=PROJECT_ID,
                                                   instance_id=INSTANCE_ID,
                                                   table_id=TABLE_ID,
                                                   task_id="id")
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with()
 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 = BigtableTableWaitForReplicationSensor(project_id=PROJECT_ID,
                                                instance_id=INSTANCE_ID,
                                                table_id=TABLE_ID,
                                                task_id="id")
     self.assertTrue(op.poke(None))
     mock_hook.assert_called_once_with()
    def test_wait_no_instance(self, mock_hook):
        mock_hook.return_value.get_instance.return_value = None

        op = BigtableTableWaitForReplicationSensor(
            project_id=PROJECT_ID,
            instance_id=INSTANCE_ID,
            table_id=TABLE_ID,
            task_id="id"
        )
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with()
    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 = BigtableTableWaitForReplicationSensor(project_id=PROJECT_ID,
                                                   instance_id=INSTANCE_ID,
                                                   table_id=TABLE_ID,
                                                   task_id="id")
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with()
 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 = BigtableTableWaitForReplicationSensor(
         project_id=PROJECT_ID,
         instance_id=INSTANCE_ID,
         table_id=TABLE_ID,
         task_id="id"
     )
     self.assertTrue(op.poke(None))
     mock_hook.assert_called_once_with()
    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 = BigtableTableWaitForReplicationSensor(
            project_id=PROJECT_ID,
            instance_id=INSTANCE_ID,
            table_id=TABLE_ID,
            task_id="id"
        )
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with()