Beispiel #1
0
    def test_updating_cluster_that_does_not_exists_empty_project_id(
            self, mock_hook):
        instance = mock_hook.return_value.get_instance.return_value = mock.Mock(
            Instance)
        mock_hook.return_value.update_cluster.side_effect = mock.Mock(
            side_effect=google.api_core.exceptions.NotFound(
                "Cluster not found."))

        with self.assertRaises(AirflowException) as e:
            op = BigtableUpdateClusterOperator(
                instance_id=INSTANCE_ID,
                cluster_id=CLUSTER_ID,
                nodes=NODES,
                task_id="id",
                gcp_conn_id=GCP_CONN_ID,
                impersonation_chain=IMPERSONATION_CHAIN,
            )
            op.execute(None)

        err = e.exception
        self.assertEqual(
            str(err),
            "Dependency: cluster '{}' does not exist for instance '{}'.".
            format(CLUSTER_ID, INSTANCE_ID),
        )
        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.update_cluster.assert_called_once_with(
            instance=instance, cluster_id=CLUSTER_ID, nodes=NODES)
Beispiel #2
0
    def test_different_error_reraised(self, mock_hook):
        op = BigtableUpdateClusterOperator(
            project_id=PROJECT_ID,
            instance_id=INSTANCE_ID,
            cluster_id=CLUSTER_ID,
            nodes=NODES,
            task_id="id",
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )
        instance = mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
        mock_hook.return_value.update_cluster.side_effect = mock.Mock(
            side_effect=google.api_core.exceptions.GoogleAPICallError('error')
        )

        with self.assertRaises(google.api_core.exceptions.GoogleAPICallError):
            op.execute(None)

        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.update_cluster.assert_called_once_with(
            instance=instance, cluster_id=CLUSTER_ID, nodes=NODES
        )
Beispiel #3
0
    def test_updating_cluster_but_instance_does_not_exists_empty_project_id(
            self, mock_hook):
        mock_hook.return_value.get_instance.return_value = None

        with self.assertRaises(AirflowException) as e:
            op = BigtableUpdateClusterOperator(instance_id=INSTANCE_ID,
                                               cluster_id=CLUSTER_ID,
                                               nodes=NODES,
                                               task_id="id",
                                               gcp_conn_id=GCP_CONN_ID)
            op.execute(None)

        err = e.exception
        self.assertEqual(
            str(err),
            "Dependency: instance '{}' does not exist.".format(INSTANCE_ID))
        mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
        mock_hook.return_value.update_cluster.assert_not_called()
Beispiel #4
0
    def test_updating_cluster_but_instance_does_not_exists_empty_project_id(
            self, mock_hook):
        mock_hook.return_value.get_instance.return_value = None

        with pytest.raises(AirflowException) as ctx:
            op = BigtableUpdateClusterOperator(
                instance_id=INSTANCE_ID,
                cluster_id=CLUSTER_ID,
                nodes=NODES,
                task_id="id",
                gcp_conn_id=GCP_CONN_ID,
                impersonation_chain=IMPERSONATION_CHAIN,
            )
            op.execute(None)

        err = ctx.value
        assert str(
            err) == f"Dependency: instance '{INSTANCE_ID}' does not exist."
        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.update_cluster.assert_not_called()