rename base sensor to match convention and add it to base classes list

apache · Apr 10, 2024 · fe2fd35 · fe2fd35
1 parent cd9fc92
commit fe2fd35
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/airflow/providers/amazon/aws/sensors/bedrock.py b/airflow/providers/amazon/aws/sensors/bedrock.py
@@ -34,7 +34,7 @@
     from airflow.utils.context import Context
 
 
-class BaseBedrockSensor(AwsBaseSensor[BedrockHook]):
+class BedrockBaseSensor(AwsBaseSensor[BedrockHook]):
     """
     General sensor behavior for Amazon Bedrock.
 
@@ -83,7 +83,7 @@ def get_state(self) -> str:
         """Implement in subclasses."""
 
 
-class BedrockCustomizeModelCompletedSensor(BaseBedrockSensor):
+class BedrockCustomizeModelCompletedSensor(BedrockBaseSensor):
     """
     Poll the state of the model customization job until it reaches a terminal state; fails if the job fails.
 
@@ -148,7 +148,7 @@ def get_state(self) -> str:
         return self.hook.conn.get_model_customization_job(jobIdentifier=self.job_name)["status"]
 
 
-class BedrockProvisionModelThroughputCompletedSensor(BaseBedrockSensor):
+class BedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor):
     """
     Poll the provisioned model throughput job until it reaches a terminal state; fails if the job fails.
 

diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py
@@ -525,6 +525,7 @@ class TestAmazonProviderProjectStructure(ExampleCoverageTest):
         "airflow.providers.amazon.aws.operators.rds.RdsBaseOperator",
         "airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator",
         "airflow.providers.amazon.aws.sensors.base_aws.AwsBaseSensor",
+        "airflow.providers.amazon.aws.sensors.bedrock.BedrockBaseSensor",
         "airflow.providers.amazon.aws.sensors.dms.DmsTaskBaseSensor",
         "airflow.providers.amazon.aws.sensors.emr.EmrBaseSensor",
         "airflow.providers.amazon.aws.sensors.rds.RdsBaseSensor",