diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py index 517199be513cb..e1e5d7144019a 100644 --- a/airflow/operators/docker_operator.py +++ b/airflow/operators/docker_operator.py @@ -206,16 +206,16 @@ def execute(self, context): self.container = self.cli.create_container( command=self.get_command(), - cpu_shares=cpu_shares, environment=self.environment, host_config=self.cli.create_host_config( binds=self.volumes, + cpu_shares=cpu_shares, + mem_limit=self.mem_limit, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, dns_search=self.dns_search), image=image, - mem_limit=self.mem_limit, user=self.user, working_dir=self.working_dir ) diff --git a/tests/operators/docker_operator.py b/tests/operators/docker_operator.py index ea90c53c28d82..d3c928ad361de 100644 --- a/tests/operators/docker_operator.py +++ b/tests/operators/docker_operator.py @@ -64,18 +64,20 @@ def test_execute(self, client_class_mock, mkdtemp_mock): client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None, version='1.19') - client_mock.create_container.assert_called_with(command='env', cpu_shares=1024, + client_mock.create_container.assert_called_with(command='env', environment={ 'AIRFLOW_TMP_DIR': '/tmp/airflow', 'UNIT': 'TEST' }, host_config=host_config, image='ubuntu:latest', - mem_limit=None, user=None, + user=None, working_dir='/container/path' ) client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path', '/mkdtemp:/tmp/airflow'], + cpu_shares=1024, + mem_limit=None, network_mode='bridge', shm_size=1000, dns=None,