def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerBaseOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=config
)
def setUp(self):
configuration.load_test_config()
from airflow.contrib.hooks.ssh_hook import SSHHook
hook = SSHHook(ssh_conn_id='ssh_default')
hook.no_host_key_check = True
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE,
'provide_context': True
}
dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
dag.schedule_interval = '@once'
self.hook = hook
self.dag = dag
self.test_dir = "/tmp"
self.test_local_dir = "/tmp/tmp2"
self.test_remote_dir = "/tmp/tmp1"
self.test_local_filename = 'test_local_file'
self.test_remote_filename = 'test_remote_file'
self.test_local_filepath = '{0}/{1}'.format(self.test_dir,
self.test_local_filename)
# Local Filepath with Intermediate Directory
self.test_local_filepath_int_dir = '{0}/{1}'.format(self.test_local_dir,
self.test_local_filename)
self.test_remote_filepath = '{0}/{1}'.format(self.test_dir,
self.test_remote_filename)
# Remote Filepath with Intermediate Directory
self.test_remote_filepath_int_dir = '{0}/{1}'.format(self.test_remote_dir,
self.test_remote_filename)
def setUp(self):
super().setUp()
self.remote_log_base = 's3://bucket/remote/log/location'
self.remote_log_location = 's3://bucket/remote/log/location/1.log'
self.remote_log_key = 'remote/log/location/1.log'
self.local_log_location = 'local/log/location'
self.filename_template = '{try_number}.log'
self.s3_task_handler = S3TaskHandler(
self.local_log_location,
self.remote_log_base,
self.filename_template
)
configuration.load_test_config()
date = datetime(2016, 1, 1)
self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
task = DummyOperator(task_id='task_for_testing_file_log_handler', dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=date)
self.ti.try_number = 1
self.ti.state = State.RUNNING
self.addCleanup(self.dag.clear)
self.conn = boto3.client('s3')
# We need to create the bucket since this is all in Moto's 'virtual'
# AWS account
moto.core.moto_api_backend.reset()
self.conn.create_bucket(Bucket="bucket")
def setUp(self):
configuration.load_test_config()
args = {
'owner': 'airflow',
'start_date': datetime.datetime(2017, 1, 1)
}
self.dag = DAG('test_dag_id', default_args=args)
def setUp(self):
configuration.load_test_config()
self.athena = AWSAthenaOperator(task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE',
database='TEST_DATABASE', output_location='s3://test_s3_bucket/',
client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595',
sleep_time=1)
def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerEndpointConfigOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_endpoint_config_params
)
def setUp(self):
super(TestLogView, self).setUp()
# Create a custom logging configuration
configuration.load_test_config()
logging_config = copy.deepcopy(DEFAULT_LOGGING_CONFIG)
current_dir = os.path.dirname(os.path.abspath(__file__))
logging_config['handlers']['task']['base_log_folder'] = os.path.normpath(
os.path.join(current_dir, 'test_logs'))
logging_config['handlers']['task']['filename_template'] = \
'{{ ti.dag_id }}/{{ ti.task_id }}/{{ ts | replace(":", ".") }}/{{ try_number }}.log'
# Write the custom logging configuration to a file
self.settings_folder = tempfile.mkdtemp()
settings_file = os.path.join(self.settings_folder, "airflow_local_settings.py")
new_logging_file = "LOGGING_CONFIG = {}".format(logging_config)
with open(settings_file, 'w') as handle:
handle.writelines(new_logging_file)
sys.path.append(self.settings_folder)
conf.set('core', 'logging_config_class', 'airflow_local_settings.LOGGING_CONFIG')
app = application.create_app(testing=True)
self.app = app.test_client()
self.session = Session()
from airflow.www.views import dagbag
dag = DAG(self.DAG_ID, start_date=self.DEFAULT_DATE)
task = DummyOperator(task_id=self.TASK_ID, dag=dag)
dagbag.bag_dag(dag, parent_dag=dag, root_dag=dag)
ti = TaskInstance(task=task, execution_date=self.DEFAULT_DATE)
ti.try_number = 1
self.session.merge(ti)
self.session.commit()
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='slack-webhook-default',
extra='{"webhook_token": "your_token_here"}')
)
def setUp(self):
configuration.load_test_config()
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE
}
self.dag = DAG(TEST_DAG_ID, default_args=args)
def setUp(self):
session = requests.Session()
adapter = requests_mock.Adapter()
session.mount('mock', adapter)
self.get_hook = HttpHook(method='GET')
self.post_hook = HttpHook(method='POST')
configuration.load_test_config()
def setUp(self):
conf.load_test_config()
app = application.create_app(testing=True)
app.config['WTF_CSRF_METHODS'] = []
self.app = app.test_client()
self.session = Session()
models.DagBag().get_dag("example_bash_operator").sync_to_db()
def setUp(self):
configuration.load_test_config()
from airflow.contrib.hooks.ssh_hook import SSHHook
from airflow.hooks.S3_hook import S3Hook
hook = SSHHook(ssh_conn_id='ssh_default')
s3_hook = S3Hook('aws_default')
hook.no_host_key_check = True
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE,
'provide_context': True
}
dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
dag.schedule_interval = '@once'
self.hook = hook
self.s3_hook = s3_hook
self.ssh_client = self.hook.get_conn()
self.sftp_client = self.ssh_client.open_sftp()
self.dag = dag
self.s3_bucket = BUCKET
self.sftp_path = SFTP_PATH
self.s3_key = S3_KEY
def setUp(self):
if sys.version_info[0] == 3:
raise unittest.SkipTest('TestSparkSubmitHook won\'t work with '
'python3. No need to test anything here')
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='spark_yarn_cluster', conn_type='spark',
host='yarn://yarn-master', extra='{"queue": "root.etl", "deploy-mode": "cluster"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_default_mesos', conn_type='spark',
host='mesos://host', port=5050)
)
db.merge_conn(
models.Connection(
conn_id='spark_home_set', conn_type='spark',
host='yarn://yarn-master',
extra='{"spark-home": "/opt/myspark"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_home_not_set', conn_type='spark',
host='yarn://yarn-master')
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='jdbc_default', conn_type='jdbc',
host='jdbc://localhost/', port=443,
extra='{"extra__jdbc__drv_path": "/path1/test.jar,/path2/t.jar2", "extra__jdbc__drv_clsname": "com.driver.main"}'))
def setUp(self):
configuration.load_test_config()
self._upload_dataframe()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.database = 'airflow'
self.table = 'hive_server_hook'
self.hql = """
CREATE DATABASE IF NOT EXISTS {{ params.database }};
USE {{ params.database }};
DROP TABLE IF EXISTS {{ params.table }};
CREATE TABLE IF NOT EXISTS {{ params.table }} (
a int,
b int)
ROW FORMAT DELIMITED
FIELDS TERMINATED BY ',';
LOAD DATA LOCAL INPATH '{{ params.csv_path }}'
OVERWRITE INTO TABLE {{ params.table }};
"""
self.columns = ['{}.a'.format(self.table),
'{}.b'.format(self.table)]
self.hook = HiveMetastoreHook()
t = HiveOperator(
task_id='HiveHook_' + str(random.randint(1, 10000)),
params={
'database': self.database,
'table': self.table,
'csv_path': self.local_path
},
hive_cli_conn_id='beeline_default',
hql=self.hql, dag=self.dag)
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
def setUp(self):
configuration.load_test_config()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.next_day = (DEFAULT_DATE +
datetime.timedelta(days=1)).isoformat()[:10]
self.database = 'airflow'
self.partition_by = 'ds'
self.table = 'static_babynames_partitioned'
self.hql = """
CREATE DATABASE IF NOT EXISTS {{ params.database }};
USE {{ params.database }};
DROP TABLE IF EXISTS {{ params.table }};
CREATE TABLE IF NOT EXISTS {{ params.table }} (
state string,
year string,
name string,
gender string,
num int)
PARTITIONED BY ({{ params.partition_by }} string);
ALTER TABLE {{ params.table }}
ADD PARTITION({{ params.partition_by }}='{{ ds }}');
"""
self.hook = HiveMetastoreHook()
t = HiveOperator(
task_id='HiveHook_' + str(random.randint(1, 10000)),
params={
'database': self.database,
'table': self.table,
'partition_by': self.partition_by
},
hive_cli_conn_id='beeline_default',
hql=self.hql, dag=self.dag)
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='cassandra_test', conn_type='cassandra',
host='host-1,host-2', port='9042', schema='test_keyspace',
extra='{"load_balancing_policy":"TokenAwarePolicy"'))
def setUp(self):
super(TestElasticsearchTaskHandler, self).setUp()
self.local_log_location = 'local/log/location'
self.filename_template = '{try_number}.log'
self.log_id_template = '{dag_id}-{task_id}-{execution_date}-{try_number}'
self.end_of_log_mark = 'end_of_log\n'
self.es_task_handler = ElasticsearchTaskHandler(
self.local_log_location,
self.filename_template,
self.log_id_template,
self.end_of_log_mark
)
self.es = elasticsearch.Elasticsearch(hosts=[{'host': 'localhost', 'port': 9200}])
self.index_name = 'test_index'
self.doc_type = 'log'
self.test_message = 'some random stuff'
self.body = {'message': self.test_message, 'log_id': self.LOG_ID,
'offset': 1}
self.es.index(index=self.index_name, doc_type=self.doc_type,
body=self.body, id=1)
configuration.load_test_config()
self.dag = DAG(self.DAG_ID, start_date=self.EXECUTION_DATE)
task = DummyOperator(task_id=self.TASK_ID, dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=self.EXECUTION_DATE)
self.ti.try_number = 1
self.ti.state = State.RUNNING
self.addCleanup(self.dag.clear)
def test_xcom_disable_pickle_type(self):
configuration.load_test_config()
json_obj = {"key": "value"}
execution_date = timezone.utcnow()
key = "xcom_test1"
dag_id = "test_dag1"
task_id = "test_task1"
configuration.set("core", "enable_xcom_pickling", "False")
XCom.set(key=key,
value=json_obj,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
ret_value = XCom.get_one(key=key,
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date)
self.assertEqual(ret_value, json_obj)
session = settings.Session()
ret_value = session.query(XCom).filter(XCom.key == key, XCom.dag_id == dag_id,
XCom.task_id == task_id,
XCom.execution_date == execution_date
).first().value
self.assertEqual(ret_value, json_obj)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='jira_default', conn_type='jira',
host='https://localhost/jira/', port=443,
extra='{"verify": "False", "project": "AIRFLOW"}'))
def setUp(self):
conf.load_test_config()
self.app, self.appbuilder = application.create_app(testing=True)
self.app.config['WTF_CSRF_ENABLED'] = False
self.client = self.app.test_client()
self.session = Session()
self.login()
def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerModelOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_model_params
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='cassandra_test', conn_type='cassandra',
host='host-1,host-2', port='9042', schema='test_keyspace',
extra='{"load_balancing_policy":"TokenAwarePolicy"}'))
db.merge_conn(
models.Connection(
conn_id='cassandra_default_with_schema', conn_type='cassandra',
host='cassandra', port='9042', schema='s'))
hook = CassandraHook("cassandra_default")
session = hook.get_conn()
cqls = [
"DROP SCHEMA IF EXISTS s",
"""
CREATE SCHEMA s WITH REPLICATION =
{ 'class' : 'SimpleStrategy', 'replication_factor' : 1 }
""",
]
for cql in cqls:
session.execute(cql)
session.shutdown()
hook.shutdown_cluster()
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='spark_yarn_cluster', conn_type='spark',
host='yarn://yarn-master',
extra='{"queue": "root.etl", "deploy-mode": "cluster"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_k8s_cluster', conn_type='spark',
host='k8s://https://k8s-master',
extra='{"spark-home": "/opt/spark", ' +
'"deploy-mode": "cluster", ' +
'"namespace": "mynamespace"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_default_mesos', conn_type='spark',
host='mesos://host', port=5050)
)
db.merge_conn(
models.Connection(
conn_id='spark_home_set', conn_type='spark',
host='yarn://yarn-master',
extra='{"spark-home": "/opt/myspark"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_home_not_set', conn_type='spark',
host='yarn://yarn-master')
)
db.merge_conn(
models.Connection(
conn_id='spark_binary_set', conn_type='spark',
host='yarn', extra='{"spark-binary": "custom-spark-submit"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_binary_and_home_set', conn_type='spark',
host='yarn',
extra='{"spark-home": "/path/to/spark_home", ' +
'"spark-binary": "custom-spark-submit"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_standalone_cluster', conn_type='spark',
host='spark://spark-standalone-master:6066',
extra='{"spark-home": "/path/to/spark_home", "deploy-mode": "cluster"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_standalone_cluster_client_mode', conn_type='spark',
host='spark://spark-standalone-master:6066',
extra='{"spark-home": "/path/to/spark_home", "deploy-mode": "client"}')
)
def setUp(self):
configuration.load_test_config()
self.sensor = AthenaSensor(task_id='test_athena_sensor',
query_execution_id='abc',
sleep_time=5,
max_retires=1,
aws_conn_id='aws_default')
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='sqoop_test', conn_type='sqoop', schema='schema',
host='rmdbs', port=5050, extra=json.dumps(self._config_json)
)
)
def setUp(self):
super(TestMountPoint, self).setUp()
configuration.load_test_config()
configuration.conf.set("webserver", "base_url", "http://localhost:8080/test")
config = dict()
config['WTF_CSRF_METHODS'] = []
app = application.cached_app(config=config, testing=True)
self.client = Client(app)
def setUp(self):
configuration.load_test_config()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
dag = DAG('test_dag_id', default_args=args)
self.dag = dag
self.sql = 'SELECT 1'
self.hook = AwsDynamoDBHook(
aws_conn_id='aws_default', region_name='us-east-1')
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='default-discord-webhook',
host='https://discordapp.com/api/',
extra='{"webhook_endpoint": "webhooks/00000/some-discord-token_000"}')
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='spark_default', conn_type='spark',
host='yarn://yarn-master')
)
def setUp(self):
configuration.load_test_config()
try:
configuration.conf.add_section("api")
except DuplicateSectionError:
pass
configuration.conf.set("api",
"auth_backend",
"airflow.contrib.auth.backends.password_auth")
self.app = application.create_app(testing=True)
session = Session()
user = models.User()
password_user = PasswordUser(user)
password_user.username = '******'
password_user.password = '******'
session.add(password_user)
session.commit()
session.close()
def setUp(self):
if sys.version_info[0] == 3:
raise unittest.SkipTest('TestSparkSubmitHook won\'t work with '
'python3. No need to test anything here')
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='spark_yarn_cluster', conn_type='spark',
host='yarn://yarn-master', extra='{"queue": "root.etl", "deploy-mode": "cluster"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_default_mesos', conn_type='spark',
host='mesos://host', port=5050)
)
db.merge_conn(
models.Connection(
conn_id='spark_home_set', conn_type='spark',
host='yarn://yarn-master',
extra='{"spark-home": "/opt/myspark"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_home_not_set', conn_type='spark',
host='yarn://yarn-master')
)
db.merge_conn(
models.Connection(
conn_id='spark_binary_set', conn_type='spark',
host='yarn', extra='{"spark-binary": "custom-spark-submit"}')
)
db.merge_conn(
models.Connection(
conn_id='spark_binary_and_home_set', conn_type='spark',
host='yarn', extra='{"spark-home": "/path/to/spark_home", "spark-binary": "custom-spark-submit"}')
)
def setUp(self):
super(TestLogView, self).setUp()
configuration.load_test_config()
logging_config = copy.deepcopy(DEFAULT_LOGGING_CONFIG)
current_dir = os.path.dirname(os.path.abspath(__file__))
logging_config['handlers']['file.task'][
'base_log_folder'] = os.path.normpath(
os.path.join(current_dir, 'test_logs'))
logging.config.dictConfig(logging_config)
app = application.create_app(testing=True)
self.app = app.test_client()
self.session = Session()
from airflow.www.views import dagbag
dag = DAG(self.DAG_ID, start_date=self.DEFAULT_DATE)
task = DummyOperator(task_id=self.TASK_ID, dag=dag)
dagbag.bag_dag(dag, parent_dag=dag, root_dag=dag)
ti = TaskInstance(task=task, execution_date=self.DEFAULT_DATE)
ti.try_number = 1
self.session.merge(ti)
self.session.commit()
def setUp(self):
super(TestElasticsearchTaskHandler, self).setUp()
self.local_log_location = 'local/log/location'
self.filename_template = '{try_number}.log'
self.log_id_template = '{dag_id}-{task_id}-{execution_date}-{try_number}'
self.end_of_log_mark = 'end_of_log\n'
self.write_stdout = False
self.json_format = False
self.json_fields = 'asctime,filename,lineno,levelname,message'
self.es_task_handler = ElasticsearchTaskHandler(
self.local_log_location, self.filename_template,
self.log_id_template, self.end_of_log_mark, self.write_stdout,
self.json_format, self.json_fields)
self.es = elasticsearch.Elasticsearch(hosts=[{
'host': 'localhost',
'port': 9200
}])
self.index_name = 'test_index'
self.doc_type = 'log'
self.test_message = 'some random stuff'
self.body = {
'message': self.test_message,
'log_id': self.LOG_ID,
'offset': 1
}
self.es.index(index=self.index_name,
doc_type=self.doc_type,
body=self.body,
id=1)
configuration.load_test_config()
self.dag = DAG(self.DAG_ID, start_date=self.EXECUTION_DATE)
task = DummyOperator(task_id=self.TASK_ID, dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=self.EXECUTION_DATE)
self.ti.try_number = 1
self.ti.state = State.RUNNING
self.addCleanup(self.dag.clear)
def setUpClass(cls):
super().setUpClass()
configuration.load_test_config()
dag_id = "extract_document_feature_prepare_dag"
cls.prepare_dag = DAG(dag_id=dag_id, start_date=DEFAULT_DATE)
get_list = GetEDINETDocumentListOperator(
task_id="get_document_list", dag=cls.prepare_dag)
get_document = GetEDINETDocumentSensor(
max_retrieve=3, document_ids=("S100E2NM","S100E2S2"),
task_id="get_document", dag=cls.prepare_dag, poke_interval=2)
register_document = RegisterDocumentOperator(
task_id="register_document", dag=cls.prepare_dag)
extract_feature = ExtractDocumentFeaturesOperator(
report_kinds=("annual",),
task_id="extract_feature", dag=cls.prepare_dag)
cls.prepare_dag.clear()
get_list.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
get_document.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
register_document.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
extract_feature.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def setUp(self):
super(TestLogView, self).setUp()
# Make sure that the configure_logging is not cached
self.old_modules = dict(sys.modules)
conf.load_test_config()
# Create a custom logging configuration
configuration.load_test_config()
logging_config = copy.deepcopy(DEFAULT_LOGGING_CONFIG)
current_dir = os.path.dirname(os.path.abspath(__file__))
logging_config['handlers']['task'][
'base_log_folder'] = os.path.normpath(
os.path.join(current_dir, 'test_logs'))
logging_config['handlers']['task']['filename_template'] = \
'{{ ti.dag_id }}/{{ ti.task_id }}/{{ ts | replace(":", ".") }}/{{ try_number }}.log'
# Write the custom logging configuration to a file
self.settings_folder = tempfile.mkdtemp()
settings_file = os.path.join(self.settings_folder,
"airflow_local_settings.py")
new_logging_file = "LOGGING_CONFIG = {}".format(logging_config)
with open(settings_file, 'w') as handle:
handle.writelines(new_logging_file)
sys.path.append(self.settings_folder)
conf.set('core', 'logging_config_class',
'airflow_local_settings.LOGGING_CONFIG')
app = application.create_app(testing=True)
self.app = app.test_client()
self.session = Session()
from airflow.www.views import dagbag
dag = DAG(self.DAG_ID, start_date=self.DEFAULT_DATE)
task = DummyOperator(task_id=self.TASK_ID, dag=dag)
dagbag.bag_dag(dag, parent_dag=dag, root_dag=dag)
ti = TaskInstance(task=task, execution_date=self.DEFAULT_DATE)
ti.try_number = 1
self.session.merge(ti)
self.session.commit()
def setUp(self):
configuration.load_test_config()
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE,
'provide_context': True
}
dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
dag.schedule_interval = '@once'
self.dag = dag
self.sensor = gcs_sensor.GoogleCloudStorageUploadSessionCompleteSensor(
task_id='sensor',
bucket='test-bucket',
prefix='test-prefix/path',
inactivity_period=12,
poke_interval=10,
min_objects=1,
allow_delete=False,
previous_num_objects=0,
dag=self.dag)
self.last_mocked_date = datetime(2019, 4, 24, 0, 0, 0)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='spark_yarn_cluster',
conn_type='spark',
host='yarn://yarn-master',
extra='{"queue": "root.etl", "deploy-mode": "cluster"}'))
db.merge_conn(
models.Connection(conn_id='spark_default_mesos',
conn_type='spark',
host='mesos://host',
port=5050))
db.merge_conn(
models.Connection(conn_id='spark_home_set',
conn_type='spark',
host='yarn://yarn-master',
extra='{"spark-home": "/opt/myspark"}'))
db.merge_conn(
models.Connection(conn_id='spark_home_not_set',
conn_type='spark',
host='yarn://yarn-master'))
db.merge_conn(
models.Connection(conn_id='spark_binary_set',
conn_type='spark',
host='yarn',
extra='{"spark-binary": "custom-spark-submit"}'))
db.merge_conn(
models.Connection(
conn_id='spark_binary_and_home_set',
conn_type='spark',
host='yarn',
extra=
'{"spark-home": "/path/to/spark_home", "spark-binary": "custom-spark-submit"}'
))
def test_templates(self, _):
dag_id = 'test_dag_id'
configuration.load_test_config()
args = {'start_date': DEFAULT_DATE}
self.dag = DAG(dag_id, default_args=args)
op = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
source_bucket='{{ dag.dag_id }}',
destination_bucket='{{ dag.dag_id }}',
description='{{ dag.dag_id }}',
object_conditions={'exclude_prefixes': ['{{ dag.dag_id }}']},
gcp_conn_id='{{ dag.dag_id }}',
task_id=TASK_ID,
dag=self.dag,
)
ti = TaskInstance(op, DEFAULT_DATE)
ti.render_templates()
self.assertEqual(dag_id, getattr(op, 'source_bucket'))
self.assertEqual(dag_id, getattr(op, 'destination_bucket'))
self.assertEqual(dag_id, getattr(op, 'description'))
self.assertEqual(
dag_id,
getattr(op, 'object_conditions')['exclude_prefixes'][0])
self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
def test_instance_start_with_templates(self, _):
dag_id = 'test_dag_id'
configuration.load_test_config()
args = {
'start_date': DEFAULT_DATE
}
self.dag = DAG(dag_id, default_args=args)
op = GceInstanceStartOperator(
project_id='{{ dag.dag_id }}',
zone='{{ dag.dag_id }}',
resource_id='{{ dag.dag_id }}',
gcp_conn_id='{{ dag.dag_id }}',
api_version='{{ dag.dag_id }}',
task_id='id',
dag=self.dag
)
ti = TaskInstance(op, DEFAULT_DATE)
ti.render_templates()
self.assertEqual(dag_id, getattr(op, 'project_id'))
self.assertEqual(dag_id, getattr(op, 'zone'))
self.assertEqual(dag_id, getattr(op, 'resource_id'))
self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
self.assertEqual(dag_id, getattr(op, 'api_version'))
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(conn_id='azure_container_instance_test',
conn_type='azure_container_instances',
login='******',
password='******',
extra=json.dumps({
'tenantId': 'tenant_id',
'subscriptionId': 'subscription_id'
})))
self.resources = ResourceRequirements(
requests=ResourceRequests(memory_in_gb='4', cpu='1'))
with patch(
'azure.common.credentials.ServicePrincipalCredentials.__init__',
autospec=True,
return_value=None):
with patch(
'azure.mgmt.containerinstance.ContainerInstanceManagementClient'
):
self.testHook = AzureContainerInstanceHook(
conn_id='azure_container_instance_test')
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='mongo_test', conn_type='mongo',
host='mongo', port='27017', schema='test'))
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE
}
self.dag = DAG('test_dag_id', default_args=args)
hook = MongoHook('mongo_test')
hook.insert_one('foo', {'bar': 'baz'})
self.sensor = MongoSensor(
task_id='test_task',
mongo_conn_id='mongo_test',
dag=self.dag,
collection='foo',
query={'bar': 'baz'}
)
def setUp(self):
load_test_config()
db.merge_conn(
models.Connection(
conn_id='google_test',
host='google',
conn_type="google_cloud_platform",
schema='refresh_token',
login='******',
password='******'
)
)
db.merge_conn(
models.Connection(
conn_id='s3_test',
conn_type='s3',
schema='test',
extra='{"aws_access_key_id": "aws_access_key_id", "aws_secret_access_key":'
' "aws_secret_access_key"}'
)
)
self.kwargs = {
'gcp_conn_id': 'google_test',
'google_api_service_name': 'test_service',
'google_api_service_version': 'v3',
'google_api_endpoint_path': 'analyticsreporting.reports.batchGet',
'google_api_endpoint_params': {},
'google_api_pagination': False,
'google_api_num_retries': 0,
'aws_conn_id': 's3_test',
's3_destination_key': 'test/google_api_to_s3_test.csv',
's3_overwrite': True,
'task_id': 'task_id',
'dag': None
}
def setUp(self):
configuration.load_test_config()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.next_day = (DEFAULT_DATE +
datetime.timedelta(days=1)).isoformat()[:10]
self.database = 'airflow'
self.partition_by = 'ds'
self.table = 'static_babynames_partitioned'
self.hql = """
CREATE DATABASE IF NOT EXISTS {{ params.database }};
USE {{ params.database }};
DROP TABLE IF EXISTS {{ params.table }};
CREATE TABLE IF NOT EXISTS {{ params.table }} (
state string,
year string,
name string,
gender string,
num int)
PARTITIONED BY ({{ params.partition_by }} string);
ALTER TABLE {{ params.table }}
ADD PARTITION({{ params.partition_by }}='{{ ds }}');
"""
self.hook = HiveMetastoreHook()
t = operators.hive_operator.HiveOperator(
task_id='HiveHook_' + str(random.randint(1, 10000)),
params={
'database': self.database,
'table': self.table,
'partition_by': self.partition_by
},
hive_cli_conn_id='beeline_default',
hql=self.hql,
dag=self.dag)
t.run(start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE,
ignore_ti_state=True)
def setUp(self):
configuration.load_test_config()
hook = SSHHook(ssh_conn_id='ssh_default')
s3_hook = S3Hook('aws_default')
hook.no_host_key_check = True
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE,
'provide_context': True
}
dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
dag.schedule_interval = '@once'
self.hook = hook
self.s3_hook = s3_hook
self.ssh_client = self.hook.get_conn()
self.sftp_client = self.ssh_client.open_sftp()
self.dag = dag
self.s3_bucket = BUCKET
self.sftp_path = SFTP_PATH
self.s3_key = S3_KEY
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(
task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1',
group='group',
placement_constraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
network_configuration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
}
)
def setUp(self):
super(TestS3TaskHandler, self).setUp()
self.remote_log_location = 's3://bucket/remote/log/location'
self.remote_log_key = 'remote/log/location'
self.local_log_location = 'local/log/location'
self.filename_template = '{try_number}.log'
self.s3_task_handler = S3TaskHandler(self.local_log_location,
self.remote_log_location,
self.filename_template)
configuration.load_test_config()
date = datetime(2016, 1, 1)
self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
task = DummyOperator(task_id='task_for_testing_file_log_handler',
dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=date)
self.ti.try_number = 1
self.addCleanup(self.dag.clear)
self.conn = boto3.client('s3')
# We need to create the bucket since this is all in Moto's 'virtual'
# AWS account
moto.core.moto_api_backend.reset()
self.conn.create_bucket(Bucket="bucket")
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs_operator_args = {
'task_id': 'task',
'task_definition': 't',
'cluster': 'c',
'overrides': {},
'aws_conn_id': None,
'region_name': 'eu-west-1',
'group': 'group',
'placement_constraints': [{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
'network_configuration': {
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
}
}
self.ecs = ECSOperator(**self.ecs_operator_args)
def setUp(self):
configuration.load_test_config()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
dag = DAG(TEST_DAG_ID, default_args=args)
self.dag = dag
def setUp(self):
super().setUp()
configuration.load_test_config()
app, _ = application.create_app(testing=True)
self.app = app.test_client()
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import print_function
from airflow import DAG, configuration, operators
from airflow.utils.tests import skipUnlessImported
from airflow.utils import timezone
import os
import mock
import unittest
configuration.load_test_config()
DEFAULT_DATE = timezone.datetime(2015, 1, 1)
DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat()
DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10]
TEST_DAG_ID = 'unit_test_dag'
@skipUnlessImported('airflow.operators.mysql_operator', 'MySqlOperator')
class MySqlTest(unittest.TestCase):
def setUp(self):
configuration.load_test_config()
args = {
'owner': 'airflow',
'mysql_conn_id': 'airflow_db',
'start_date': DEFAULT_DATE
def setUpClass(cls):
os.environ['AIRFLOW__TESTSECTION__TESTKEY'] = 'testvalue'
os.environ['AIRFLOW__TESTSECTION__TESTPERCENT'] = 'with%percent'
configuration.load_test_config()
conf.set('core', 'percent', 'with%%inside')
def airflow_init_db(airflow_home):
configuration.load_test_config()
initdb()
def setUp(self):
configuration.load_test_config()
self.channel_mock = mock.patch('grpc.Channel').start()
def setUp(self):
configuration.load_test_config()
def setUp(self):
configuration.load_test_config()
self.s3_test_url = "s3://test/this/is/not/a-real-key.txt"
def setUp(self):
configuration.load_test_config()
self.hook = SFTPHook()
os.makedirs(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
with open(os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS), 'a') as f:
f.write('Test file')
def setUp(self):
super(TestSnowflakeOperator, self).setUp()
configuration.load_test_config()
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
dag = DAG(TEST_DAG_ID, default_args=args)
self.dag = dag
def setUp(self):
configuration.load_test_config()
self.dagbag = DagBag(dag_folder=DEV_NULL, include_examples=True)
self.args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
dag = DAG(TEST_DAG_ID, default_args=self.args)
self.dag = dag
def setUp(self):
super(TestApiExperimental, self).setUp()
configuration.load_test_config()
app = application.create_app(testing=True)
self.app = app.test_client()
