def test_create_task(self, mock_hook):
mock_hook.return_value.create_task.return_value = TEST_TASK
operator = CloudTasksTaskCreateOperator(location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
task_id="id")
result = operator.execute(context=None)
self.assertEqual({'appEngineHttpRequest': {}}, result)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
)
mock_hook.return_value.create_task.assert_called_once_with(
location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
project_id=None,
task_name=None,
response_view=None,
retry=None,
timeout=None,
metadata=None,
)
def test_create_task(self, mock_hook):
mock_hook.return_value.create_task.return_value = {}
operator = CloudTasksTaskCreateOperator(
location=LOCATION, queue_name=QUEUE_ID, task=Task(), task_id="id"
)
operator.execute(context=None)
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
mock_hook.return_value.create_task.assert_called_once_with(
location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
project_id=None,
task_name=None,
response_view=None,
retry=None,
timeout=None,
metadata=None,
)
def test_create_task(self, get_conn):
result = self.hook.create_task(
location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
project_id=PROJECT_ID,
task_name=TASK_NAME,
)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.create_task.assert_called_once_with(
parent=FULL_QUEUE_PATH,
task=Task(name=FULL_TASK_PATH),
response_view=None,
retry=None,
timeout=None,
metadata=None,
)
def test_create_task(self, mock_hook):
mock_hook.return_value.create_task.return_value = TEST_TASK
operator = CloudTasksTaskCreateOperator(location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
task_id="id")
result = operator.execute(context=None)
self.assertEqual(
{
'app_engine_http_request': {
'body': '',
'headers': {},
'http_method': 0,
'relative_uri': ''
},
'dispatch_count': 0,
'name': '',
'response_count': 0,
'view': 0,
},
result,
)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
)
mock_hook.return_value.create_task.assert_called_once_with(
location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
project_id=None,
task_name=None,
response_view=None,
retry=None,
timeout=None,
metadata=None,
)
def test_list_tasks(self, get_conn):
result = self.hook.list_tasks(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, [Task(name=FULL_TASK_PATH)])
get_conn.return_value.list_tasks.assert_called_once_with(
request=dict(parent=FULL_QUEUE_PATH,
response_view=None,
page_size=None),
retry=None,
timeout=None,
metadata=(),
)
def execute(self, context: 'Context'):
hook = CloudTasksHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
task = hook.get_task(
location=self.location,
queue_name=self.queue_name,
task_name=self.task_name,
project_id=self.project_id,
response_view=self.response_view,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Task.to_dict(task)
def execute(self, context: 'Context'):
hook = CloudTasksHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
tasks = hook.list_tasks(
location=self.location,
queue_name=self.queue_name,
project_id=self.project_id,
response_view=self.response_view,
page_size=self.page_size,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return [Task.to_dict(t) for t in tasks]
def test_run_task(self, get_conn):
result = self.hook.run_task(
location=LOCATION,
queue_name=QUEUE_ID,
task_name=TASK_NAME,
project_id=PROJECT_ID,
)
self.assertEqual(result, Task(name=FULL_TASK_PATH))
get_conn.return_value.run_task.assert_called_once_with(
request=dict(name=FULL_TASK_PATH, response_view=None),
retry=None,
timeout=None,
metadata=(),
)
def setUp(self):
super(HandlerTest, self).setUp()
self.handler = util.Handler(self.request, self.response)
FakeGrSource.clear()
util.now_fn = lambda: NOW
# add FakeSource everywhere necessary
util.BLACKLIST.add('fa.ke')
util.webmention_endpoint_cache.clear()
self.stubbed_create_task = False
tasks_client.create_task = lambda *args, **kwargs: Task(name='foo')
# clear datastore
orig_requests_post('http://%s/reset' % ndb_client.host)
self.ndb_context = ndb_client.context()
self.ndb_context.__enter__()
def expect_task(self, queue, eta_seconds=None, **kwargs):
if not self.stubbed_create_task:
self.mox.StubOutWithMock(tasks_client, 'create_task')
self.stubbed_create_task = True
def check_queue(path):
if not path.endswith('/' + queue):
# print("expect_task: %s doesn't end with /%s!" % (path, queue))
return False
return True
def check_params(params):
req = params['app_engine_http_request']
if not check_queue(req['relative_uri']):
return False
# convert model objects and keys to url-safe key strings for comparison
for name, val in kwargs.items():
if isinstance(val, ndb.Model):
kwargs[name] = val.key.urlsafe().decode()
elif isinstance(val, ndb.Key):
kwargs[name] = val.urlsafe().decode()
got = set(urllib.parse.parse_qsl(req['body'].decode()))
expected = set(kwargs.items())
if got != expected:
# print('expect_task: expected %s, got %s' % (expected, got))
return False
if eta_seconds is not None:
got = params['schedule_time'].seconds - util.to_utc_timestamp(
util.now_fn())
delta = eta_seconds * .2 + 10
if not (got + delta >= eta_seconds >= got - delta):
# print('expect_task: expected schedule_time %r, got %r' % (eta_seconds, got))
return False
return True
return tasks_client.create_task(
mox.Func(check_queue),
mox.Func(check_params)).InAnyOrder().AndReturn(
Task(name='my task'))
def expect_task(self, queue, eta_seconds=None, **kwargs):
self.stub_create_task()
def check_task(task):
if not task.parent.endswith('/' + queue):
# These can help for debugging, but can also be misleading, since many
# tests insert multiple tasks, so check_task() runs on all of them (due
# to InAnyOrder() below) until it finds one that matches.
# print("expect_task: %s doesn't end with /%s!" % (task.parent, queue))
return False
req = task.task.app_engine_http_request
if not req.relative_uri.endswith('/' + queue):
# print("expect_task: relative_uri %s doesn't end with /%s!" % (
# req.relative_uri, queue))
return False
# convert model objects and keys to url-safe key strings for comparison
for name, val in kwargs.items():
if isinstance(val, ndb.Model):
kwargs[name] = val.key.urlsafe().decode()
elif isinstance(val, ndb.Key):
kwargs[name] = val.urlsafe().decode()
got = set(urllib.parse.parse_qsl(req.body.decode()))
expected = set(kwargs.items())
if got != expected:
# print('expect_task: expected %s, got %s' % (expected, got))
return False
if eta_seconds is not None:
got = (util.to_utc_timestamp(task.task.schedule_time) -
util.to_utc_timestamp(util.now_fn()))
delta = eta_seconds * .2 + 10
if not (got + delta >= eta_seconds >= got - delta):
# print('expect_task: expected schedule_time %r, got %r' % (eta_seconds, got))
return False
return True
return tasks_client.create_task(
mox.Func(check_task)).InAnyOrder().AndReturn(Task(name='my task'))
def execute(self, context: 'Context'):
hook = CloudTasksHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
task = hook.run_task(
location=self.location,
queue_name=self.queue_name,
task_name=self.task_name,
project_id=self.project_id,
response_view=self.response_view,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
CloudTasksQueueLink.persist(
operator_instance=self,
context=context,
queue_name=task.name,
)
return Task.to_dict(task)
def execute(self, context: 'Context'):
hook = CloudTasksHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
tasks = hook.list_tasks(
location=self.location,
queue_name=self.queue_name,
project_id=self.project_id,
response_view=self.response_view,
page_size=self.page_size,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
CloudTasksQueueLink.persist(
operator_instance=self,
context=context,
queue_name=f"projects/{self.project_id or hook.project_id}/"
f"locations/{self.location}/queues/{self.queue_name}",
)
return [Task.to_dict(t) for t in tasks]
CloudTasksTaskDeleteOperator,
CloudTasksTaskGetOperator,
CloudTasksTaskRunOperator,
CloudTasksTasksListOperator,
)
GCP_CONN_ID = "google_cloud_default"
PROJECT_ID = "test-project"
LOCATION = "asia-east2"
FULL_LOCATION_PATH = "projects/test-project/locations/asia-east2"
QUEUE_ID = "test-queue"
FULL_QUEUE_PATH = "projects/test-project/locations/asia-east2/queues/test-queue"
TASK_NAME = "test-task"
FULL_TASK_PATH = "projects/test-project/locations/asia-east2/queues/test-queue/tasks/test-task"
TEST_QUEUE = Queue(name=FULL_QUEUE_PATH)
TEST_TASK = Task(app_engine_http_request={})
class TestCloudTasksQueueCreate(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook"
)
def test_create_queue(self, mock_hook):
mock_hook.return_value.create_queue.return_value = TEST_QUEUE
operator = CloudTasksQueueCreateOperator(location=LOCATION,
task_queue=TEST_QUEUE,
task_id="id")
result = operator.execute(context=None)
self.assertEqual({'name': FULL_QUEUE_PATH, 'state': 0}, result)
mock_hook.assert_called_once_with(
def setUp(self):
super().setUp()
FakeGrSource.clear()
util.now_fn = lambda: NOW
# add FakeSource everywhere necessary
util.BLOCKLIST.add('fa.ke')
util.webmention_endpoint_cache.clear()
self.stubbed_create_task = False
tasks_client.create_task = lambda *args, **kwargs: Task(name='foo')
self.client = self.app.test_client()
self.client.__enter__()
self.clear_datastore()
self.ndb_context = ndb_client.context()
self.ndb_context.__enter__()
# sources
self.auth_entities = [
FakeAuthEntity(key=ndb.Key('FakeAuthEntity', '01122334455'),
user_json=json_dumps({
'id': '0123456789',
'name': 'Fake User',
'url': 'http://fakeuser.com/',
})),
FakeAuthEntity(key=ndb.Key('FakeAuthEntity', '0022446688'),
user_json=json_dumps({
'id':
'0022446688',
'name':
'Another Fake',
'url':
'http://anotherfake.com/',
}))
]
self.sources = [
FakeSource.new(auth_entity=self.auth_entities[0]),
FakeSource.new(auth_entity=self.auth_entities[1])
]
for entity in self.sources:
entity.features = ['listen']
with self.app.test_request_context():
self.source_bridgy_url = self.sources[0].bridgy_url()
self.actor = FakeGrSource.actor = {
'objectType': 'person',
'id': 'tag:fa.ke,2013:212038',
'username': '******',
'displayName': 'Ryan B',
'url': 'https://snarfed.org/',
'image': {
'url': 'http://pic.ture/url'
},
}
# activities
self.activities = FakeGrSource.activities = [{
'id': f'tag:source.com,2013:{id}',
'url': 'http://fa.ke/post/url',
'object': {
'objectType':
'note',
'id':
f'tag:source.com,2013:{id}',
'url':
'http://fa.ke/post/url',
'content':
'foo http://target1/post/url bar',
'to': [{
'objectType': 'group',
'alias': '@public'
}],
'replies': {
'items': [{
'objectType': 'comment',
'id': f'tag:source.com,2013:1_2_{id}',
'url': 'http://fa.ke/comment/url',
'content': 'foo bar',
}],
'totalItems':
1,
},
'tags': [{
'objectType': 'activity',
'verb': 'like',
'id': f'tag:source.com,2013:{id}_liked_by_alice',
'object': {
'url': 'http://example.com/abc'
},
'author': {
'id': 'tag:source.com,2013:alice',
'url': 'http://example.com/alice',
},
}, {
'id': f'tag:source.com,2013:{id}_reposted_by_bob',
'objectType': 'activity',
'verb': 'share',
'object': {
'url': 'http://example.com/def'
},
'author': {
'url': 'http://example.com/bob'
},
}, {
'id': f'tag:source.com,2013:{id}_scissors_by_bob',
'objectType': 'activity',
'verb': 'react',
'content': '✁',
'object': {
'url': 'http://example.com/def'
},
'author': {
'url': 'http://example.com/bob'
},
}],
},
} for id in ('a', 'b', 'c')]
# responses
self.responses = []
created = datetime.now(timezone.utc) - timedelta(days=10)
for activity in self.activities:
obj = activity['object']
pruned_activity = {
'id': activity['id'],
'url': 'http://fa.ke/post/url',
'object': {
'content': 'foo http://target1/post/url bar',
}
}
comment = obj['replies']['items'][0]
self.responses.append(
Response(id=comment['id'],
activities_json=[json_dumps(pruned_activity)],
response_json=json_dumps(comment),
type='comment',
source=self.sources[0].key,
unsent=['http://target1/post/url'],
created=created))
created += timedelta(hours=1)
like = obj['tags'][0]
self.responses.append(
Response(id=like['id'],
activities_json=[json_dumps(pruned_activity)],
response_json=json_dumps(like),
type='like',
source=self.sources[0].key,
unsent=['http://target1/post/url'],
created=created))
created += timedelta(hours=1)
share = obj['tags'][1]
self.responses.append(
Response(id=share['id'],
activities_json=[json_dumps(pruned_activity)],
response_json=json_dumps(share),
type='repost',
source=self.sources[0].key,
unsent=['http://target1/post/url'],
created=created))
created += timedelta(hours=1)
reaction = obj['tags'][2]
self.responses.append(
Response(id=reaction['id'],
activities_json=[json_dumps(pruned_activity)],
response_json=json_dumps(reaction),
type='react',
source=self.sources[0].key,
unsent=['http://target1/post/url'],
created=created))
created += timedelta(hours=1)
# publishes
self.publishes = [
Publish(
parent=PublishedPage(id='https://post').key,
source=self.sources[0].key,
status='complete',
published={'url': 'http://fa.ke/syndpost'},
)
]
# blogposts
self.blogposts = [
BlogPost(
id='https://post',
source=self.sources[0].key,
status='complete',
feed_item={'title': 'a post'},
sent=['http://a/link'],
)
]
class TestCloudTasksHook(unittest.TestCase):
def setUp(self):
with mock.patch(
"airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudTasksHook(gcp_conn_id="test")
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.client_info",
new_callable=mock.PropertyMock,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook._get_credentials"
)
@mock.patch("airflow.providers.google.cloud.hooks.tasks.CloudTasksClient")
def test_cloud_tasks_client_creation(self, mock_client, mock_get_creds,
mock_client_info):
result = self.hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.hook._client, result)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.create_queue.return_value":
API_RESPONSE}, # type: ignore
)
def test_create_queue(self, get_conn):
result = self.hook.create_queue(
location=LOCATION,
task_queue=Queue(),
queue_name=QUEUE_ID,
project_id=PROJECT_ID,
)
self.assertIs(result, API_RESPONSE)
get_conn.return_value.create_queue.assert_called_once_with(
request=dict(parent=FULL_LOCATION_PATH,
queue=Queue(name=FULL_QUEUE_PATH)),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.update_queue.return_value":
API_RESPONSE}, # type: ignore
)
def test_update_queue(self, get_conn):
result = self.hook.update_queue(
task_queue=Queue(state=3),
location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID,
)
self.assertIs(result, API_RESPONSE)
get_conn.return_value.update_queue.assert_called_once_with(
request=dict(queue=Queue(name=FULL_QUEUE_PATH, state=3),
update_mask=None),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.get_queue.return_value":
API_RESPONSE}, # type: ignore
)
def test_get_queue(self, get_conn):
result = self.hook.get_queue(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertIs(result, API_RESPONSE)
get_conn.return_value.get_queue.assert_called_once_with(
request=dict(name=FULL_QUEUE_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{
"return_value.list_queues.return_value":
[Queue(name=FULL_QUEUE_PATH)]
}, # type: ignore
)
def test_list_queues(self, get_conn):
result = self.hook.list_queues(location=LOCATION,
project_id=PROJECT_ID)
self.assertEqual(result, [Queue(name=FULL_QUEUE_PATH)])
get_conn.return_value.list_queues.assert_called_once_with(
request=dict(parent=FULL_LOCATION_PATH,
filter=None,
page_size=None),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.delete_queue.return_value": None}, # type: ignore
)
def test_delete_queue(self, get_conn):
result = self.hook.delete_queue(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, None)
get_conn.return_value.delete_queue.assert_called_once_with(
request=dict(name=FULL_QUEUE_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{
"return_value.purge_queue.return_value":
Queue(name=FULL_QUEUE_PATH)
}, # type: ignore
)
def test_purge_queue(self, get_conn):
result = self.hook.purge_queue(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))
get_conn.return_value.purge_queue.assert_called_once_with(
request=dict(name=FULL_QUEUE_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{
"return_value.pause_queue.return_value":
Queue(name=FULL_QUEUE_PATH)
}, # type: ignore
)
def test_pause_queue(self, get_conn):
result = self.hook.pause_queue(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))
get_conn.return_value.pause_queue.assert_called_once_with(
request=dict(name=FULL_QUEUE_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{
"return_value.resume_queue.return_value":
Queue(name=FULL_QUEUE_PATH)
}, # type: ignore
)
def test_resume_queue(self, get_conn):
result = self.hook.resume_queue(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))
get_conn.return_value.resume_queue.assert_called_once_with(
request=dict(name=FULL_QUEUE_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.create_task.return_value":
Task(name=FULL_TASK_PATH)}, # type: ignore
)
def test_create_task(self, get_conn):
result = self.hook.create_task(
location=LOCATION,
queue_name=QUEUE_ID,
task=Task(),
project_id=PROJECT_ID,
task_name=TASK_NAME,
)
self.assertEqual(result, Task(name=FULL_TASK_PATH))
get_conn.return_value.create_task.assert_called_once_with(
request=dict(parent=FULL_QUEUE_PATH,
task=Task(name=FULL_TASK_PATH),
response_view=None),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.get_task.return_value":
Task(name=FULL_TASK_PATH)}, # type: ignore
)
def test_get_task(self, get_conn):
result = self.hook.get_task(
location=LOCATION,
queue_name=QUEUE_ID,
task_name=TASK_NAME,
project_id=PROJECT_ID,
)
self.assertEqual(result, Task(name=FULL_TASK_PATH))
get_conn.return_value.get_task.assert_called_once_with(
request=dict(name=FULL_TASK_PATH, response_view=None),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{
"return_value.list_tasks.return_value":
[Task(name=FULL_TASK_PATH)]
}, # type: ignore
)
def test_list_tasks(self, get_conn):
result = self.hook.list_tasks(location=LOCATION,
queue_name=QUEUE_ID,
project_id=PROJECT_ID)
self.assertEqual(result, [Task(name=FULL_TASK_PATH)])
get_conn.return_value.list_tasks.assert_called_once_with(
request=dict(parent=FULL_QUEUE_PATH,
response_view=None,
page_size=None),
retry=None,
timeout=None,
metadata=(),
)
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.delete_task.return_value": None}, # type: ignore
)
def test_delete_task(self, get_conn):
result = self.hook.delete_task(
location=LOCATION,
queue_name=QUEUE_ID,
task_name=TASK_NAME,
project_id=PROJECT_ID,
)
self.assertEqual(result, None)
get_conn.return_value.delete_task.assert_called_once_with(
request=dict(name=FULL_TASK_PATH),
retry=None,
timeout=None,
metadata=())
@mock.patch(
"airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
**{"return_value.run_task.return_value":
Task(name=FULL_TASK_PATH)}, # type: ignore
)
def test_run_task(self, get_conn):
result = self.hook.run_task(
location=LOCATION,
queue_name=QUEUE_ID,
task_name=TASK_NAME,
project_id=PROJECT_ID,
)
self.assertEqual(result, Task(name=FULL_TASK_PATH))
get_conn.return_value.run_task.assert_called_once_with(
request=dict(name=FULL_TASK_PATH, response_view=None),
retry=None,
timeout=None,
metadata=(),
)
