Taragolis commented on code in PR #35226:
URL: https://github.com/apache/airflow/pull/35226#discussion_r1376095456


##########
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:
   I have no idea why use literals always a better that call constructor, that 
is like personal preference over the strict rule. The only one exception which 
I know is: Empty dict literal `{}` should use over `dict()` and this one 
detected by `black` or/and `ruff`.
   
   In particular this dictionary use for provide attributes to class 
constructor. So personally I found that in that case use dict constructor allow 
contributors to easy copy/paste from Operators/Sensors constructor attributes 
to default test arguments and vice versa without adding remove `"` to attribute 
names and replace `:` -> `=` -> `:`



-- 
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

Reply via email to