From f93d1dc3ef76715fce44b6ae781a8460a2633ab6 Mon Sep 17 00:00:00 2001
From: jrybicki-jsc <j.rybicki@fz-juelich.de>
Date: Mon, 13 Sep 2021 11:58:00 +0200
Subject: [PATCH] fixing tests
---
dags/b2shareoperator.py | 7 +++---
requirements.txt | 4 +++-
tests/test_b2shareoperator.py | 41 +++++++++++++++++++++++++----------
3 files changed, 36 insertions(+), 16 deletions(-)
diff --git a/dags/b2shareoperator.py b/dags/b2shareoperator.py
index 1824dac..ac76449 100644
--- a/dags/b2shareoperator.py
+++ b/dags/b2shareoperator.py
@@ -40,18 +40,17 @@ class B2ShareOperator(BaseOperator):
self.name = name
self.conn_id = conn_id
self.target_dir = target_dir
- print(self.target_dir)
-
+
def execute(self, **kwargs):
hook = HttpHook(http_conn_id=self.conn_id, method='GET')
- print(kwargs)
params = kwargs['context']['params']
oid = params['oid']
+
hrespo = hook.run(endpoint=f"/api/records/{oid}")
print(hrespo)
flist = get_file_list(hrespo.json())
- print(flist)
+
ti = kwargs['context']['ti']
name_mappings = {}
for fname, url in flist.items():
diff --git a/requirements.txt b/requirements.txt
index 1f88981..240d00f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,5 @@
requests
urllib3==1.26.6
-apache-airflow-providers-ssh
\ No newline at end of file
+apache-airflow-providers-ssh
+apache-airflow-providers-http
+apache-airflow-providers-sftp
diff --git a/tests/test_b2shareoperator.py b/tests/test_b2shareoperator.py
index 5d56e6a..751bca6 100644
--- a/tests/test_b2shareoperator.py
+++ b/tests/test_b2shareoperator.py
@@ -1,17 +1,21 @@
import unittest
-from airflow.utils.state import State
-from airflow.utils.dates import days_ago
-from dags.b2shareoperator import B2ShareOperator
+from unittest.mock import patch, Mock
+
from airflow import DAG
from airflow.models.taskinstance import TaskInstance
+from airflow.utils.dates import days_ago
+from airflow.utils.state import State
+from dags.b2shareoperator import B2ShareOperator, get_file_list
DEFAULT_DATE = '2019-10-03'
TEST_DAG_ID = 'test_my_custom_operator'
+
class B2ShareOperatorTest(unittest.TestCase):
- def setUp(self):
- self.dag = DAG(TEST_DAG_ID, schedule_interval='@daily', default_args={'start_date' : days_ago(2)}, params={"oid": "111"})
+ def setUp(self):
+ self.dag = DAG(TEST_DAG_ID, schedule_interval='@daily',
+ default_args={'start_date': days_ago(2)}, params={"oid": "111"})
self.op = B2ShareOperator(
dag=self.dag,
task_id='test',
@@ -19,9 +23,24 @@ class B2ShareOperatorTest(unittest.TestCase):
)
self.ti = TaskInstance(task=self.op, execution_date=days_ago(1))
- def test_execute_no_trigger(self):
- #with conf_vars({('params': {'oid': 1111})}):
- self.ti.run(ignore_ti_state=False)
- print(self.ti.state)
- self.assertEqual(State.SUCCESS, self.ti.state)
- # Assert something related to tasks results
\ No newline at end of file
+
+ @patch('dags.b2shareoperator.HttpHook')
+ @patch('dags.b2shareoperator.get_file_list')
+ @patch('dags.b2shareoperator.download_file')
+ def test_alt_execute_no_trigger(self, down, gfl, ht):
+ gfl.return_value = {'ooo.txt': 'htt://file/to/download'}
+ down.return_value = 'tmp_name'
+
+ self.ti.run(ignore_ti_state=False)
+ print(self.ti.state)
+ self.assertEqual(State.SUCCESS, self.ti.state)
+ # Assert something related to tasks results
+
+ def test_get_files(self):
+ with patch('dags.b2shareoperator.requests.get') as get:
+ m = Mock()
+ m.json.return_value = {'contents': [{'key': 'veryimportant.txt', 'links':{'self': 'http://foo.bar'}}]}
+ get.return_value = m
+ ret = get_file_list(obj={'links': {'files': ['bla']}})
+ self.assertEqual(len(ret), 1)
+
--
GitLab