Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/tests/providers/celery/sensors/test_celery_queue.py: 14%

1# 

2# Licensed to the Apache Software Foundation (ASF) under one 

3# or more contributor license agreements. See the NOTICE file 

4# distributed with this work for additional information 

5# regarding copyright ownership. The ASF licenses this file 

6# to you under the Apache License, Version 2.0 (the 

7# "License"); you may not use this file except in compliance 

8# with the License. You may obtain a copy of the License at 

9# 

10# http://www.apache.org/licenses/LICENSE-2.0 

11# 

12# Unless required by applicable law or agreed to in writing, 

13# software distributed under the License is distributed on an 

14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

15# KIND, either express or implied. See the License for the 

16# specific language governing permissions and limitations 

17# under the License. 

18from __future__ import annotations 

19 

20from unittest.mock import patch 

21 

22from airflow.providers.celery.sensors.celery_queue import CeleryQueueSensor 

23 

24 

25class TestCeleryQueueSensor: 

26 def setup_method(self): 

27 class TestCeleryqueueSensor(CeleryQueueSensor): 

28 def _check_task_id(self, context): 

29 return True 

30 

31 self.sensor = TestCeleryqueueSensor 

32 

33 @patch("celery.app.control.Inspect") 

34 def test_poke_success(self, mock_inspect): 

35 mock_inspect_result = mock_inspect.return_value 

36 # test success 

37 mock_inspect_result.reserved.return_value = {"test_queue": []} 

38 

39 mock_inspect_result.scheduled.return_value = {"test_queue": []} 

40 

41 mock_inspect_result.active.return_value = {"test_queue": []} 

42 test_sensor = self.sensor(celery_queue="test_queue", task_id="test-task") 

43 assert test_sensor.poke(None) 

44 

45 @patch("celery.app.control.Inspect") 

46 def test_poke_fail(self, mock_inspect): 

47 mock_inspect_result = mock_inspect.return_value 

48 # test success 

49 mock_inspect_result.reserved.return_value = {"test_queue": []} 

50 

51 mock_inspect_result.scheduled.return_value = {"test_queue": []} 

52 

53 mock_inspect_result.active.return_value = {"test_queue": ["task"]} 

54 test_sensor = self.sensor(celery_queue="test_queue", task_id="test-task") 

55 assert not test_sensor.poke(None) 

56 

57 @patch("celery.app.control.Inspect") 

58 def test_poke_success_with_taskid(self, mock_inspect): 

59 test_sensor = self.sensor( 

60 celery_queue="test_queue", task_id="test-task", target_task_id="target-task" 

61 ) 

62 assert test_sensor.poke(None)