class TestGKEHookClient(unittest.TestCase):
def setUp(self):
self.gke_hook = GKEHook(location=GKE_ZONE)
@mock.patch("airflow.gcp.hooks.kubernetes_engine.GKEHook.client_info",
new_callable=mock.PropertyMock)
@mock.patch("airflow.gcp.hooks.kubernetes_engine.GKEHook._get_credentials")
@mock.patch(
"airflow.gcp.hooks.kubernetes_engine.container_v1.ClusterManagerClient"
)
def test_gke_cluster_client_creation(self, mock_client, mock_get_creds,
mock_client_info):
result = self.gke_hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.gke_hook._client, result)
class TestGKEHook(unittest.TestCase):
def setUp(self):
self.gke_hook = GKEHook(location=GKE_ZONE)
self.gke_hook._client = mock.Mock()
@mock.patch('airflow.gcp.hooks.kubernetes_engine.container_v1.'
'ClusterManagerClient')
@mock.patch('airflow.gcp.hooks.base.ClientInfo')
@mock.patch('airflow.gcp.hooks.kubernetes_engine.GKEHook._get_credentials')
def test_get_client(self, mock_get_credentials, mock_client_info,
mock_client):
self.gke_hook._client = None
self.gke_hook.get_conn()
assert mock_get_credentials.called
mock_client.assert_called_once_with(
credentials=mock_get_credentials.return_value,
client_info=mock_client_info.return_value)
def test_get_operation(self):
self.gke_hook._client.get_operation = mock.Mock()
self.gke_hook.get_operation('TEST_OP', project_id=TEST_GCP_PROJECT_ID)
self.gke_hook._client.get_operation.assert_called_once_with(
project_id=TEST_GCP_PROJECT_ID,
zone=GKE_ZONE,
operation_id='TEST_OP')
def test_append_label(self):
key = 'test-key'
val = 'test-val'
mock_proto = mock.Mock()
self.gke_hook._append_label(mock_proto, key, val)
mock_proto.resource_labels.update.assert_called_once_with({key: val})
def test_append_label_replace(self):
key = 'test-key'
val = 'test.val+this'
mock_proto = mock.Mock()
self.gke_hook._append_label(mock_proto, key, val)
mock_proto.resource_labels.update.assert_called_once_with(
{key: 'test-val-this'})
@mock.patch("airflow.gcp.hooks.kubernetes_engine.time.sleep")
def test_wait_for_response_done(self, time_mock):
from google.cloud.container_v1.gapic.enums import Operation
mock_op = mock.Mock()
mock_op.status = Operation.Status.DONE
self.gke_hook.wait_for_operation(mock_op)
self.assertEqual(time_mock.call_count, 1)
@mock.patch("airflow.gcp.hooks.kubernetes_engine.time.sleep")
def test_wait_for_response_exception(self, time_mock):
from google.cloud.container_v1.gapic.enums import Operation
from google.cloud.exceptions import GoogleCloudError
mock_op = mock.Mock()
mock_op.status = Operation.Status.ABORTING
with self.assertRaises(GoogleCloudError):
self.gke_hook.wait_for_operation(mock_op)
self.assertEqual(time_mock.call_count, 1)
@mock.patch("airflow.gcp.hooks.kubernetes_engine.GKEHook.get_operation")
@mock.patch("airflow.gcp.hooks.kubernetes_engine.time.sleep")
def test_wait_for_response_running(self, time_mock, operation_mock):
from google.cloud.container_v1.gapic.enums import Operation
running_op, done_op, pending_op = mock.Mock(), mock.Mock(), mock.Mock()
running_op.status = Operation.Status.RUNNING
done_op.status = Operation.Status.DONE
pending_op.status = Operation.Status.PENDING
# Status goes from Running -> Pending -> Done
operation_mock.side_effect = [pending_op, done_op]
self.gke_hook.wait_for_operation(running_op,
project_id=TEST_GCP_PROJECT_ID)
self.assertEqual(time_mock.call_count, 3)
operation_mock.assert_any_call(running_op.name,
project_id=TEST_GCP_PROJECT_ID)
operation_mock.assert_any_call(pending_op.name,
project_id=TEST_GCP_PROJECT_ID)
self.assertEqual(operation_mock.call_count, 2)