def test_python_operator():
dag = DAG(dag_id='foo', start_date=datetime.now())
doit = PythonOperator(task_id='doit',
provide_context=True,
python_callable=python_method,
dag=dag)
ti = TaskInstance(task=doit, execution_date=datetime.now())
doit.execute(ti.get_template_context())
def test_sort_data(test_dag):
task = PythonOperator(task_id='sort_data',
python_callable=sort_data,
dag=test_dag)
test_dag.clear()
ti = TaskInstance(task=task, execution_date=datetime.now())
task.execute(ti.get_template_context())
assert "data_sorted.csv" in os.listdir('tmp')
def test_preprocess_data(test_dag):
task = PythonOperator(task_id='preprocess_data',
python_callable=preprocess_data,
op_kwargs={'path_cats': 'Data/US_category_id.json'},
dag=test_dag)
test_dag.clear()
ti = TaskInstance(task=task, execution_date=datetime.now())
task.execute(ti.get_template_context())
assert "data_cats.csv" in os.listdir('tmp')
def test_filter_data(test_dag):
task = PythonOperator(
task_id='filter_data',
python_callable=filter_data,
op_kwargs={'path_vids': 'Data/USvideos.csv.zip', 'date': '17.14.11'},
dag=test_dag)
test_dag.clear()
ti = TaskInstance(task=task, execution_date=datetime.now())
task.execute(ti.get_template_context())
assert "data_vids.csv" in os.listdir('tmp')
def test_store_data(test_dag):
task = PythonOperator(
task_id='store_data',
python_callable=store_data,
op_kwargs={'path_output': 'Data/top_10_trendy_cats.csv'},
dag=test_dag)
test_dag.clear()
ti = TaskInstance(task=task, execution_date=datetime.now())
task.execute(ti.get_template_context())
assert "top_10_trendy_cats.csv" in os.listdir('Data')
def run_tasks(**kwargs):
# select users from BQ table for which the events are not loaded yet
query = '''
select distinct_id
from mixpanel.users
where events_loaded = False
'''
credentials = service_account.Credentials.from_service_account_file(
GCLOUD_JSON_KEY)
df = pd.read_gbq(query,
project_id=GCLOUD_PROJECT_ID,
credentials=credentials)
distinct_ids_max = 100
dids = df['distinct_id'].to_list()
# split distinct_ids list into sublists for using it in threads
chunks_dids = [
dids[x:x + distinct_ids_max]
for x in range(0, len(dids), distinct_ids_max)
]
logging.warning(len(chunks_dids))
for job_number in range(len(chunks_dids)):
upload_events_to_bq = PythonOperator(
task_id='upload_events_bq_{}'.format(str(job_number)),
python_callable=get_events,
provide_context=True,
op_kwargs={
"start": last_upload,
"end": last_date,
"num": chunks_dids[job_number]
},
dag=dag)
change_events = PythonOperator(
task_id='change_flag_{}'.format(str(job_number)),
python_callable=change_events_fl,
provide_context=True,
op_kwargs={"num": chunks_dids[job_number]},
dag=dag)
upload_events_to_bq.execute(kwargs)
change_events.execute(kwargs)
upload_events_to_bq >> change_events
Variable.set("last_update_events", last_date)
def test_android_runner_task_install_android_apks(self):
device_serial = '192.168.31.223:5555'
app_id = 'com.chi.ssetest'
release_version = 'release-20191016-0.0.3'
with DAG(dag_id='any_dag', start_date=datetime.now()) as dag:
def push_function(**kwargs):
release_apk = ReleaseFile()
release_apk.name = 'app-debug.apk'
release_apk.type = 'application/vnd.android.package-archive'
release_apk.url = 'https://github.com/stocksdktest/AndroidTestRunner/releases/download/release-20191016-0.0.3/app-debug.apk'
release_testing = ReleaseFile()
release_testing.name = 'app-debug-androidTest.apk'
release_testing.type = 'application/vnd.android.package-archive'
release_testing.url = 'https://github.com/stocksdktest/AndroidTestRunner/releases/download/release-20191016-0.0.3/app-debug-androidTest.apk'
kwargs['ti'].xcom_push(key='android_release',
value=[release_apk, release_testing])
release_provider = PythonOperator(task_id='push_task',
python_callable=push_function,
provide_context=True)
android_runner = AndroidRunnerOperator(task_id='android_runner',
provide_context=True,
apk_id=app_id,
apk_version=release_version,
target_device=device_serial,
runner_conf=RunnerConfig())
release_provider >> android_runner
execution_date = datetime.now()
provider_instance = TaskInstance(task=release_provider,
execution_date=execution_date)
release_provider.execute(provider_instance.get_template_context())
runner_instance = TaskInstance(task=android_runner,
execution_date=execution_date)
context = runner_instance.get_template_context()
context['run_id'] = 'fake-run'
android_runner.pre_execute(context)
main_apk_version = get_app_version(device_serial, app_id)
testing_apk_version = get_app_version(device_serial,
'%s.test' % app_id)
self.assertEqual(release_version, main_apk_version,
'main app version not match')
self.assertEqual(release_version, testing_apk_version,
'testing app version not match')
def test_execute(self):
dag = DAG(dag_id='example', start_date=datetime.now())
task = PythonOperator(python_callable=ret_hello,
task_id='test_python',
dag=dag)
ti = TaskInstance(task=task, execution_date=datetime.now())
result = task.execute(ti.get_template_context())
self.assertEqual(result, 'Hello')
def test_pull_parameters_by_task(env, create_dag, create_df_task, push_xcom):
dag = create_dag
# df_task = create_df_task(dag)
push_task = PythonOperator(task_id='three',
python_callable=push_xcom,
op_args=['three_value'],
dag=dag,
provide_context=True)
push_ti = TaskInstance(task=push_task, execution_date=datetime.now())
push_task.execute(push_ti.get_template_context())
xcom = XCom.get_many(execution_date=push_ti.execution_date,
dag_ids=[dag.dag_id],
task_ids=[push_task.task_id])
pprint(xcom[0])
assert len(xcom) == 1
assert xcom[0].value == 'three_value'
assert False
def test_pull_parameters_by_specific_key(env, create_dag, create_df_task,
push_xcom_key):
dag = create_dag
df_task = create_df_task(dag)
push_task = PythonOperator(task_id='two',
python_callable=push_xcom_key,
op_args=['two', 'two_value'],
dag=dag,
provide_context=True)
push_ti = TaskInstance(task=push_task, execution_date=datetime.now())
push_task.execute(push_ti.get_template_context())
xcom = XCom.get_many(execution_date=push_ti.execution_date,
dag_ids=[dag.dag_id])
assert len(xcom) == 1
assert xcom[0].value == 'two_value'
merged = df_task.merge_parameters(push_ti)
assert merged['two-specific-key'] == 'two_value'
def test_python_operator():
test = PythonOperator(task_id="test", python_callable=lambda: "testme")
result = test.execute(context={})
assert result == "testme"
def test_python_operator_with_context():
test = PythonOperator(task_id="test", python_callable=next_week, provide_context=True)
testdate = datetime.datetime(2020, 1, 1)
result = test.execute(context={"execution_date": testdate})
assert result == testdate + datetime.timedelta(days=7)
