diff --git a/cli/tests/fixtures/task.py b/cli/tests/fixtures/task.py index 02ecb12..7934a9a 100644 --- a/cli/tests/fixtures/task.py +++ b/cli/tests/fixtures/task.py @@ -1,6 +1,4 @@ -from dcos.mesos import Slave, Task - -import mock +from dcos.mesos import Task def task_fixture(): @@ -43,9 +41,6 @@ def task_fixture(): ] }, None) - task.user = mock.Mock(return_value='root') - slave = Slave({"hostname": "mock-hostname"}, None, None) - task.slave = mock.Mock(return_value=slave) return task diff --git a/cli/tests/unit/test_tables.py b/cli/tests/unit/test_tables.py index 6b5b77e..0e0bdf4 100644 --- a/cli/tests/unit/test_tables.py +++ b/cli/tests/unit/test_tables.py @@ -1,5 +1,6 @@ import datetime +from dcos.mesos import Slave from dcoscli import tables import mock @@ -14,8 +15,12 @@ from ..fixtures.task import browse_fixture, task_fixture def test_task_table(): + task = task_fixture() + task.user = mock.Mock(return_value='root') + slave = Slave({"hostname": "mock-hostname"}, None, None) + task.slave = mock.Mock(return_value=slave) _test_table(tables.task_table, - [task_fixture()], + [task], 'tests/unit/data/task.txt')