o-nikolas commented on code in PR #35226:
URL: https://github.com/apache/airflow/pull/35226#discussion_r1378245972
##########
tests/providers/amazon/aws/sensors/test_batch.py:
##########
@@ -29,59 +30,86 @@
)
from airflow.providers.amazon.aws.triggers.batch import BatchJobTrigger
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+
+
TASK_ID = "batch_job_sensor"
JOB_ID = "8222a1c2-b246-4e19-b1b8-0039bb4407c0"
AWS_REGION = "eu-west-1"
ENVIRONMENT_NAME = "environment_name"
JOB_QUEUE = "job_queue"
-
-@pytest.fixture(scope="module")
-def batch_sensor() -> BatchSensor:
- return BatchSensor(
- task_id="batch_job_sensor",
- job_id=JOB_ID,
- )
+SOFT_FAIL_CASES = [
+ pytest.param(False, AirflowException, id="not-soft-fail"),
+ pytest.param(True, AirflowSkipException, id="soft-fail"),
+]
@pytest.fixture(scope="module")
def deferrable_batch_sensor() -> BatchSensor:
return BatchSensor(task_id="task", job_id=JOB_ID, region_name=AWS_REGION,
deferrable=True)
-class TestBatchSensor:
- @mock.patch.object(BatchClientHook, "get_job_description")
- def test_poke_on_success_state(self, mock_get_job_description,
batch_sensor: BatchSensor):
- mock_get_job_description.return_value = {"status": "SUCCEEDED"}
- assert batch_sensor.poke({}) is True
- mock_get_job_description.assert_called_once_with(JOB_ID)
+@pytest.fixture
+def mock_get_job_description():
+ with mock.patch.object(BatchClientHook, "get_job_description") as m:
+ yield m
- @mock.patch.object(BatchClientHook, "get_job_description")
- def test_poke_on_failure_state(self, mock_get_job_description,
batch_sensor: BatchSensor):
- mock_get_job_description.return_value = {"status": "FAILED"}
- with pytest.raises(AirflowException, match="Batch sensor failed. AWS
Batch job status: FAILED"):
- batch_sensor.poke({})
- mock_get_job_description.assert_called_once_with(JOB_ID)
+@pytest.fixture
+def mock_batch_client():
+ with mock.patch.object(BatchClientHook, "client") as m:
+ yield m
+
+
+class BaseBatchSensorsTests:
+ """Base test class for Batch Sensors."""
+
+ op_class: type[AwsBaseSensor]
+ default_op_kwargs: dict[str, Any]
+
+ def test_base_aws_op_attributes(self):
+ op = self.op_class(**self.default_op_kwargs)
+ assert op.hook.aws_conn_id == "aws_default"
+ assert op.hook._region_name is None
+ assert op.hook._verify is None
+ assert op.hook._config is None
+
+ op = self.op_class(
+ **self.default_op_kwargs,
+ aws_conn_id="aws-test-custom-conn",
+ region_name="eu-west-1",
+ verify=False,
+ botocore_config={"read_timeout": 42},
+ )
+ assert op.hook.aws_conn_id == "aws-test-custom-conn"
+ assert op.hook._region_name == "eu-west-1"
+ assert op.hook._verify is False
+ assert op.hook._config is not None
+ assert op.hook._config.read_timeout == 42
- @mock.patch.object(BatchClientHook, "get_job_description")
- def test_poke_on_invalid_state(self, mock_get_job_description,
batch_sensor: BatchSensor):
- mock_get_job_description.return_value = {"status": "INVALID"}
- with pytest.raises(
- AirflowException, match="Batch sensor failed. Unknown AWS Batch
job status: INVALID"
- ):
- batch_sensor.poke({})
+class TestBatchSensor(BaseBatchSensorsTests):
+ op_class = BatchSensor
+
+ @pytest.fixture(autouse=True)
+ def setup_test_cases(self):
+ self.default_op_kwargs = dict(
+ task_id="test_batch_sensor",
+ job_id=JOB_ID,
+ )
Review Comment:
The literal is slightly faster I believe (not the most compelling reason in
this case, but worth noting since we're talking about it) but the more
important thing is the keys are quoted in the literal. That way you don't have
to worry about keywords or special characters in the key name. Which isn't a
worry in this case, but I think it's best to choose one and stay consistent
across the code base. So I'd also vote to move this to the literal `{...}`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
