Remove unittest.main() from tests (apache#8454)

tests/api/common/experimental/test_delete_dag.py

@@ -132,7 +132,3 @@ def test_delete_subdag_successful_delete(self):
         self.check_dag_models_exists()
         delete_dag(dag_id=self.key, keep_records_in_log=False)
         self.check_dag_models_removed(expect_logs=0)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/api/common/experimental/test_mark_tasks.py

@@ -593,7 +593,3 @@ def tearDown(self):
         with create_session() as session:
             session.query(models.DagRun).delete()
             session.query(models.TaskInstance).delete()
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/api/common/experimental/test_pool.py

@@ -140,7 +140,3 @@ def test_delete_default_pool_not_allowed(self):
         with self.assertRaisesRegex(AirflowBadRequest,
                                     "^default_pool cannot be deleted$"):
             pool_api.delete_pool(Pool.DEFAULT_POOL_NAME)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/api/common/experimental/test_trigger_dag.py

@@ -166,7 +166,3 @@ def test_trigger_dag_with_dict_conf(self, dag_bag_mock):
             replace_microseconds=True)
 
         self.assertEqual(triggers[0].conf, conf)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/executors/test_celery_executor.py

@@ -189,7 +189,3 @@ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock
 
 def test_operation_timeout_config():
     assert celery_executor.OPERATION_TIMEOUT == 2
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/executors/test_kubernetes_executor.py

@@ -302,7 +302,3 @@ def test_change_state_failed_pod_deletion(self, mock_delete_pod, mock_get_kube_c
         executor._change_state(key, State.FAILED, 'pod_id', 'test-namespace')
         self.assertTrue(executor.event_buffer[key] == State.FAILED)
         mock_delete_pod.assert_called_once_with('pod_id', 'test-namespace')
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/executors/test_local_executor.py

@@ -77,7 +77,3 @@ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock
                  mock.call('executor.queued_tasks', mock.ANY),
                  mock.call('executor.running_tasks', mock.ANY)]
         mock_stats_gauge.assert_has_calls(calls)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/executors/test_sequential_executor.py

@@ -34,7 +34,3 @@ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock
                  mock.call('executor.queued_tasks', mock.ANY),
                  mock.call('executor.running_tasks', mock.ANY)]
         mock_stats_gauge.assert_has_calls(calls)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/macros/test_hive.py

@@ -39,7 +39,3 @@ def test_closest_ds_partition(self):
         self.assertEqual("2017-04-26", str(hive._closest_date(target_dt, [date1, date2, date3, date5], None)))
         self.assertEqual("2017-04-28", str(hive._closest_date(target_dt, [date1, date2, date4, date5])))
         self.assertEqual("2017-04-26", str(hive._closest_date(target_dt, date_list)))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_aws_dynamodb.py

@@ -74,7 +74,3 @@ def test_insert_batch_items_dynamodb_table(self):
         table.meta.client.get_waiter(
             'table_exists').wait(TableName='test_airflow')
         self.assertEqual(table.item_count, 10)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_base_aws.py

@@ -208,7 +208,3 @@ def test_expand_role(self):
         arn = hook.expand_role('test-role')
         expect_arn = conn.get_role(RoleName='test-role').get('Role').get('Arn')
         self.assertEqual(arn, expect_arn)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_batch_client.py

@@ -322,7 +322,3 @@ def test_exponential_delay(self, tries, lower, upper):
         result = self.batch_client.exponential_delay(tries)
         self.assertGreaterEqual(result, lower)
         self.assertLessEqual(result, upper)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/providers/amazon/aws/hooks/test_batch_waiters.py

@@ -459,7 +459,3 @@ def test_wait_for_job_raises_for_waiter_error(self):
             self.assertEqual(get_waiter.call_args_list, [mock.call("JobExists")])
             mock_waiter.wait.assert_called_with(jobs=[self.job_id])
             self.assertEqual(mock_waiter.wait.call_count, 1)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/providers/amazon/aws/hooks/test_cloud_formation.py

@@ -98,7 +98,3 @@ def test_delete_stack(self):
         stacks = self.hook.get_conn().describe_stacks()['Stacks']
         matching_stacks = [x for x in stacks if x['StackName'] == stack_name]
         self.assertEqual(len(matching_stacks), 0, f'stack with name {stack_name} should not exist')
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_datasync.py

@@ -423,7 +423,3 @@ def test_wait_for_task_execution_timeout(self, mock_get_conn):
                 task_execution_arn, max_iterations=1
             )
             self.assertIsNone(result)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/providers/amazon/aws/hooks/test_emr.py

@@ -88,7 +88,3 @@ def test_get_cluster_id_by_name(self):
         no_match = hook.get_cluster_id_by_name('foo', ['RUNNING', 'WAITING', 'BOOTSTRAPPING'])
 
         self.assertIsNone(no_match)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_glue_catalog.py

@@ -171,7 +171,3 @@ def test_get_table_location(self):
 
         result = self.hook.get_table_location(DB_NAME, TABLE_NAME)
         self.assertEqual(result, TABLE_INPUT['StorageDescriptor']['Location'])
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_kinesis.py

@@ -67,7 +67,3 @@ def test_insert_batch_records_kinesis_firehose(self):
 
         self.assertEqual(response['FailedPutCount'], 0)
         self.assertEqual(response['ResponseMetadata']['HTTPStatusCode'], 200)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_lambda_function.py

@@ -78,7 +78,3 @@ def test_invoke_lambda_function(self):
         response = hook.invoke_lambda(payload=json.dumps(payload))
 
         self.assertEqual(response["StatusCode"], 202)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_logs.py

@@ -75,7 +75,3 @@ def test_get_log_events(self):
         assert count == 1
         assert events[0]['timestamp'] == input_events[0]['timestamp']
         assert events[0]['message'] == input_events[0]['message']
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_redshift.py

@@ -104,7 +104,3 @@ def test_cluster_status_returns_available_cluster(self):
         hook = RedshiftHook(aws_conn_id='aws_default')
         status = hook.cluster_status('test_cluster')
         self.assertEqual(status, 'available')
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_sagemaker.py

@@ -668,7 +668,3 @@ def test_training_with_logs(self, mock_describe, mock_client, mock_log_client, m
                                  print_log=True, check_interval=1)
         self.assertEqual(mock_describe.call_count, 3)
         self.assertEqual(mock_session.describe_training_job.call_count, 1)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_sns.py

@@ -76,7 +76,3 @@ def test_publish_to_target_plain(self):
         response = hook.publish_to_target(target, message)
 
         assert 'MessageId' in response
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/hooks/test_sqs.py

@@ -34,7 +34,3 @@ class TestAwsSQSHook(unittest.TestCase):
     def test_get_conn(self):
         hook = SQSHook(aws_conn_id='aws_default')
         self.assertIsNotNone(hook.get_conn())
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_athena.py

@@ -143,7 +143,3 @@ def test_xcom_push_and_pull(self, mock_conn, mock_run_query, mock_check_query_st
         self.assertEqual(ti.xcom_pull(task_ids='test_aws_athena_operator'),
                          ATHENA_QUERY_ID)
 # pylint: enable=unused-argument
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_batch.py

@@ -157,7 +157,3 @@ def test_kill_job(self):
         self.client_mock.terminate_job.assert_called_once_with(
             jobId=JOB_ID, reason="Task killed by the user"
         )
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/providers/amazon/aws/operators/test_cloud_formation.py

@@ -99,7 +99,3 @@ def test_delete_stack(self):
             operator.execute(self.mock_context)
 
         self.cloudformation_client_mock.delete_stack.assert_any_call(StackName=stack_name)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_ecs.py

@@ -304,7 +304,3 @@ def test_check_success_task_not_raises(self):
         self.ecs._check_success_task()
         client_mock.describe_tasks.assert_called_once_with(
             cluster='c', tasks=['arn'])
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_emr_add_steps.py

@@ -195,7 +195,3 @@ def test_init_with_nonexistent_cluster_name(self):
             with self.assertRaises(AirflowException) as error:
                 operator.execute(self.mock_context)
             self.assertEqual(str(error.exception), f'No cluster found for name: {cluster_name}')
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_emr_create_job_flow.py

@@ -110,7 +110,3 @@ def test_execute_returns_job_id(self):
 
         with patch('boto3.session.Session', boto3_session_mock):
             self.assertEqual(self.operator.execute(None), 'j-8989898989')
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_emr_modify_cluster.py

@@ -81,7 +81,3 @@ def test_execute_returns_error(self):
 
         with patch('boto3.session.Session', self.boto3_session_mock):
             self.assertRaises(AirflowException, self.operator.execute, self.mock_context)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_emr_terminate_job_flow.py

@@ -49,7 +49,3 @@ def test_execute_terminates_the_job_flow_and_does_not_error(self):
             )
 
             operator.execute(None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_google_api_to_s3_transfer.py

@@ -165,7 +165,3 @@ def test_execute_with_xcom_exceeded_max_xcom_size(
         )
         context['task_instance'].xcom_push.assert_not_called()
         mock_sys_getsizeof.assert_called_once_with(mock_google_api_hook_query.return_value)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_hive_to_dynamodb.py

@@ -137,7 +137,3 @@ def test_pre_process_records_with_schema(self, mock_get_pandas_df):
         table = self.hook.get_conn().Table('test_airflow')
         table.meta.client.get_waiter('table_exists').wait(TableName='test_airflow')
         self.assertEqual(table.item_count, 1)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_mongo_to_s3.py

@@ -110,7 +110,3 @@ def test_execute(self, mock_s3_hook, mock_mongo_hook):
             bucket_name=S3_BUCKET,
             replace=False
         )
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_s3_list.py

@@ -43,7 +43,3 @@ def test_execute(self, mock_hook):
         mock_hook.return_value.list_keys.assert_called_once_with(
             bucket_name=BUCKET, prefix=PREFIX, delimiter=DELIMITER)
         self.assertEqual(sorted(files), sorted(MOCK_FILES))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_s3_to_sftp.py

@@ -145,7 +145,3 @@ def delete_remote_resource(self):
 
     def tearDown(self):
         self.delete_remote_resource()
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_base.py

@@ -67,7 +67,3 @@ def test_parse_integer(self):
         ]
         self.sagemaker.parse_config_integers()
         self.assertEqual(self.sagemaker.config, parsed_config)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_endpoint.py

@@ -128,7 +128,3 @@ def test_execute_with_duplicate_endpoint_creation(self, mock_endpoint_update,
                                              'ResponseMetadata':
                                              {'HTTPStatusCode': 200}}
         self.sagemaker.execute(None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_endpoint_config.py

@@ -77,7 +77,3 @@ def test_execute_with_failure(self, mock_model, mock_client):
             }
         }
         self.assertRaises(AirflowException, self.sagemaker.execute, None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_model.py

@@ -68,7 +68,3 @@ def test_execute_with_failure(self, mock_model, mock_client):
                                    'ResponseMetadata':
                                        {'HTTPStatusCode': 404}}
         self.assertRaises(AirflowException, self.sagemaker.execute, None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_training.py

@@ -154,7 +154,3 @@ def test_execute_with_existing_job_fail(
         mock_create_training_job.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}}
         mock_list_training_jobs.return_value = [{"TrainingJobName": job_name}]
         self.assertRaises(AirflowException, self.sagemaker.execute, None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_transform.py

@@ -121,7 +121,3 @@ def test_execute_with_failure(self, mock_transform, mock_model, mock_client):
                                        'ResponseMetadata':
                                        {'HTTPStatusCode': 404}}
         self.assertRaises(AirflowException, self.sagemaker.execute, None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sagemaker_tuning.py

@@ -153,7 +153,3 @@ def test_execute_with_failure(self, mock_tuning, mock_client):
                                     'ResponseMetadata':
                                     {'HTTPStatusCode': 404}}
         self.assertRaises(AirflowException, self.sagemaker.execute, None)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sftp_to_s3.py

@@ -122,7 +122,3 @@ def test_sftp_to_s3_operation(self):
         conn.delete_object(Bucket=self.s3_bucket, Key=self.s3_key)
         conn.delete_bucket(Bucket=self.s3_bucket)
         self.assertFalse(self.s3_hook.check_for_bucket(self.s3_bucket))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sns.py

@@ -73,7 +73,3 @@ def test_execute(self, mock_hook):
 
         # Then
         self.assertEqual(hook_response, result)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/operators/test_sqs.py

@@ -67,7 +67,3 @@ def test_execute_success(self):
         context_calls = []
 
         self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context call  should be same")
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/sensors/test_athena.py

@@ -57,7 +57,3 @@ def test_poke_cancelled(self, mock_poll_query_status):
         with self.assertRaises(AirflowException) as context:
             self.sensor.poke(None)
         self.assertIn('Athena sensor failed', str(context.exception))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/sensors/test_cloud_formation.py

@@ -120,7 +120,3 @@ def test_poke_stack_in_unsuccessful_state(self):
     def test_poke_stack_does_not_exist(self):
         op = CloudFormationDeleteStackSensor(task_id='task', stack_name='foo')
         self.assertTrue(op.poke({}))
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/providers/amazon/aws/sensors/test_emr_base.py

@@ -110,7 +110,3 @@ def test_poke_raises_error_when_state_is_in_failed_states(self):
         self.assertIn('EMR job failed', str(context.exception))
         self.assertIn(EXPECTED_CODE, str(context.exception))
         self.assertNotIn(EMPTY_CODE, str(context.exception))
-
-
-if __name__ == '__main__':
-    unittest.main()