airflow/providers/docker/operators/docker: extra_hosts support added

airflow/providers/docker/operators/docker.py

@@ -167,6 +167,7 @@ def __init__(
         shm_size: Optional[int] = None,
         tty: Optional[bool] = False,
         cap_add: Optional[Iterable[str]] = None,
+        extra_hosts: Optional[Dict[str, str]] = None,
         **kwargs,
     ) -> None:
 
@@ -200,6 +201,7 @@ def __init__(
         self.shm_size = shm_size
         self.tty = tty
         self.cap_add = cap_add
+        self.extra_hosts = extra_hosts
         if kwargs.get('xcom_push') is not None:
             raise AirflowException("'xcom_push' was deprecated, use 'BaseOperator.do_xcom_push' instead")
 
@@ -244,6 +246,7 @@ def _run_image(self) -> Optional[str]:
                     cpu_shares=int(round(self.cpus * 1024)),
                     mem_limit=self.mem_limit,
                     cap_add=self.cap_add,
+                    extra_hosts=self.extra_hosts,
                 ),
                 image=self.image,
                 user=self.user,

tests/providers/docker/operators/test_docker.py

@@ -246,3 +246,18 @@ def test_execute_xcom_behavior(self):
 
         self.assertEqual(xcom_push_result, b'container log')
         self.assertIs(no_xcom_push_result, None)
+
+    def test_extra_hosts(self):
+        hosts_obj = mock.Mock()
+        operator = DockerOperator(task_id='test', image='test', extra_hosts=hosts_obj)
+        operator.execute(None)
+        self.client_mock.create_container.assert_called_once()
+        self.assertIn(
+            'host_config', self.client_mock.create_container.call_args.kwargs,
+        )
+        self.assertIn(
+            'extra_hosts', self.client_mock.create_host_config.call_args.kwargs,
+        )
+        self.assertIs(
+            hosts_obj, self.client_mock.create_host_config.call_args.kwargs['extra_hosts'],
+        )