def testRunMissingDesiredPod(self, exec_mock, tmp_kubeconfig_mock):
test_env_object = self.MakeEnvironmentWithStateAndClusterLocation(
self.messages.Environment.StateValueValuesEnum.RUNNING)
self.ExpectEnvironmentGet(self.TEST_PROJECT,
self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID,
response=test_env_object)
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeFetchKubectlNamespaceCallback([('default', 'Active')]))
fake_exec.AddCallback(
1,
self.MakeGetPodsCallback([
command_util.GkePodStatus('pod1', 'running', 'true'),
command_util.GkePodStatus('pod2', 'running', 'true')
]))
with self.AssertRaisesExceptionMatches(command_util.Error,
'Desired GKE pod not found'):
self.RunEnvironments('run', '--project', self.TEST_PROJECT,
'--location', self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID,
self.TEST_SUBCOMMAND)
fake_exec.Verify()
def testGetGkePod_PodNoReady(self, exec_mock, tmp_kubeconfig_mock):
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeGetPodsCallback(
[command_util.GkePodStatus('airflow-worker1', 'creating', 'true'),
command_util.GkePodStatus('airflow-worker2', 'running', '')]))
pod = command_util.GetGkePod('airflow-worker',
self.TEST_KUBECTL_DEFAULT_NAMESPACE)
self.assertEqual('airflow-worker2', pod)
fake_exec.Verify()
def testGetGkePod_PodNotRunning(self, exec_mock, tmp_kubeconfig_mock):
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeGetPodsCallback(
[command_util.GkePodStatus('pod1', 'creating', 'false'),
command_util.GkePodStatus('pod2', 'creating', 'false')]))
with self.AssertRaisesExceptionMatches(command_util.Error,
'No running GKE pods found.'):
command_util.GetGkePod('pod1', self.TEST_KUBECTL_DEFAULT_NAMESPACE)
fake_exec.Verify()
def testRunWithArgs(self, exec_mock, tmp_kubeconfig_mock):
test_env_object = self.MakeEnvironmentWithStateAndClusterLocation(
self.messages.Environment.StateValueValuesEnum.RUNNING)
self.ExpectEnvironmentGet(self.TEST_PROJECT,
self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID,
response=test_env_object)
subcmd_args = ['--arg1', '--arg2', 'val2']
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeFetchKubectlNamespaceCallback([('default', 'Active')]))
fake_exec.AddCallback(
1,
self.MakeGetPodsCallback(
[command_util.GkePodStatus(self.TEST_POD, 'running', 'true')]))
fake_exec.AddCallback(
2,
self.MakeKubectlExecCallback(self.TEST_POD,
subcmd_args=subcmd_args))
self.RunEnvironments('run', '--project', self.TEST_PROJECT,
'--location', self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID, self.TEST_SUBCOMMAND,
'--', *subcmd_args)
fake_exec.Verify()
def testSubcommandRun_BackfillOlderAirflow_YesNotAdded(
self, exec_mock, tmp_kubeconfig_mock):
test_env_object = self.MakeEnvironmentWithStateAndClusterLocation(
self.messages.Environment.StateValueValuesEnum.RUNNING,
image_version='composer-1.2.3-airflow-1.10.3')
self.ExpectEnvironmentGet(self.TEST_PROJECT,
self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID,
response=test_env_object)
subcmd = 'backfill'
subcmd_args = []
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeFetchKubectlNamespaceCallback([('default', 'Active')]))
fake_exec.AddCallback(
1,
self.MakeGetPodsCallback(
[command_util.GkePodStatus(self.TEST_POD, 'running', 'true')]))
# Ensure that the '--yes' argument has not been added to the list of
# 'backfill' args.
fake_exec.AddCallback(
2,
self.MakeKubectlExecCallback(self.TEST_POD,
subcmd=subcmd,
subcmd_args=subcmd_args))
self.RunEnvironments('run', '--project', self.TEST_PROJECT,
'--location', self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID, subcmd, '--',
*subcmd_args)
fake_exec.Verify()
def testSubcommandRun_DeleteDag(self, exec_mock, tmp_kubeconfig_mock):
test_env_object = self.MakeEnvironmentWithStateAndClusterLocation(
self.messages.Environment.StateValueValuesEnum.RUNNING)
self.ExpectEnvironmentGet(self.TEST_PROJECT,
self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID,
response=test_env_object)
subcmd = 'delete_dag'
subcmd_args = []
fake_exec = kubectl_util.FakeExec()
exec_mock.side_effect = fake_exec
tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig
fake_exec.AddCallback(
0, self.MakeFetchKubectlNamespaceCallback([('default', 'Active')]))
fake_exec.AddCallback(
1,
self.MakeGetPodsCallback(
[command_util.GkePodStatus(self.TEST_POD, 'running', 'true')]))
# Ensure that the '--yes' argument is added to the list of 'delete_dag' args
# if it is not present.
fake_exec.AddCallback(
2,
self.MakeKubectlExecCallback(self.TEST_POD,
subcmd=subcmd,
subcmd_args=subcmd_args + ['--yes']))
self.RunEnvironments('run', '--project', self.TEST_PROJECT,
'--location', self.TEST_LOCATION,
self.TEST_ENVIRONMENT_ID, subcmd, '--',
*subcmd_args)
fake_exec.Verify()