Python asynctest.CoroutineMock() Examples

def test_service_can_be_unregistered_by_instance():
    with mock.patch('bookshelf.discolib.client') as client_func:
        client = asynctest.CoroutineMock()
        svc = mock.MagicMock()
        
        client_func.return_value = client
        client.agent.services.register.return_value = svc
        
        res = await discolib.register_service(id="svc1", name="svc",
                                              tags=["a", "b"],
                                              port=9000,
                                              address="localhost")
        assert res == svc

        client.agent.services.deregister.return_value = True
        res = await discolib.deregister_service(svc)
        client.agent.services.deregister.assert_called_with(svc)
        assert res is True 

def test_callhandlers_calls_handlers_for_loglevel(self):
        level10_handler = Mock(level=10, handle=CoroutineMock())
        level30_handler = Mock(level=30, handle=CoroutineMock())

        logger = Logger.with_default_handlers()
        logger.handlers = [level10_handler, level30_handler]

        record = LogRecord(
            level=20,
            name="aiologger",
            pathname="/aiologger/tests/test_logger.py",
            lineno=17,
            msg="Xablau!",
            exc_info=None,
            args=None,
        )
        await logger.call_handlers(record)

        level10_handler.handle.assert_awaited_once_with(record)
        level30_handler.handle.assert_not_awaited() 

def test_it_calls_multiple_handlers_if_multiple_handle_matches_are_found_for_record(
        self
    ):
        level10_handler = Mock(level=10, handle=CoroutineMock())
        level20_handler = Mock(level=20, handle=CoroutineMock())

        logger = Logger.with_default_handlers()
        logger.handlers = [level10_handler, level20_handler]

        record = LogRecord(
            level=30,
            name="aiologger",
            pathname="/aiologger/tests/test_logger.py",
            lineno=17,
            msg="Xablau!",
            exc_info=None,
            args=None,
        )

        await logger.call_handlers(record)

        level10_handler.handle.assert_awaited_once_with(record)
        level20_handler.handle.assert_awaited_once_with(record) 

def test_emit_awaits_for_handle_error_is_an_exceptions_is_raised(
        self
    ):
        handler = BaseAsyncRotatingFileHandler(filename=self.temp_file.name)
        handler.should_rollover = Mock(return_value=False)
        exc = OSError()
        with patch(
            "aiologger.handlers.files.AsyncFileHandler.emit", side_effect=exc
        ), patch.object(
            handler, "handle_error", CoroutineMock()
        ) as handleError:
            log_record = LogRecord(
                name="Xablau",
                level=20,
                pathname="/aiologger/tests/test_logger.py",
                lineno=17,
                msg="Xablau!",
                exc_info=None,
                args=None,
            )
            await handler.emit(log_record)
            handleError.assert_awaited_once_with(log_record, exc) 

def test_get_all_running_tasks():
    with asynctest.patch(
        "paasta_tools.mesos_tools.get_current_tasks", autospec=True
    ) as mock_get_current_tasks, asynctest.patch(
        "paasta_tools.mesos_tools.filter_running_tasks", autospec=True
    ) as mock_filter_running_tasks, asynctest.patch(
        "paasta_tools.mesos_tools.get_mesos_master", autospec=True
    ) as mock_get_mesos_master:
        mock_task_1 = mock.Mock()
        mock_task_2 = mock.Mock()
        mock_task_3 = mock.Mock()

        mock_get_current_tasks.return_value = [mock_task_1, mock_task_2]
        mock_orphan_tasks = asynctest.CoroutineMock(return_value=[mock_task_3])
        mock_mesos_master = mock.Mock(orphan_tasks=mock_orphan_tasks)
        mock_get_mesos_master.return_value = mock_mesos_master

        ret = await mesos_tools.get_all_running_tasks()
        mock_get_current_tasks.assert_called_with("")
        mock_filter_running_tasks.assert_called_with(
            [mock_task_1, mock_task_2, mock_task_3]
        )
        assert ret == mock_filter_running_tasks.return_value 

def test_check_mesos_active_frameworks_fails(capfd):
    with mock.patch(
        "paasta_tools.monitoring.check_mesos_active_frameworks.parse_args",
        autospec=True,
    ) as mock_parse_args, mock.patch(
        "paasta_tools.monitoring.check_mesos_active_frameworks.get_mesos_master",
        autospec=True,
    ) as mock_get_mesos_master:
        mock_opts = mock.MagicMock()
        mock_opts.expected = "foo,bar"
        mock_parse_args.return_value = mock_opts
        mock_master = mock.MagicMock()
        mock_master.state = CoroutineMock(
            func=CoroutineMock(),  # https://github.com/notion/a_sync/pull/40
            return_value={"frameworks": [{"name": "foo"}]},
        )
        mock_get_mesos_master.return_value = mock_master

        with pytest.raises(SystemExit) as error:
            check_mesos_active_frameworks()
        out, err = capfd.readouterr()
        assert "CRITICAL" in out
        assert "bar" in out
        assert "foo" not in out
        assert error.value.code == 2 

def test_emit_calls_handle_error_if_an_error_occurs(self):
        writer = Mock(write=CoroutineMock(), drain=CoroutineMock())
        with patch(
            "aiologger.handlers.streams.StreamWriter", return_value=writer
        ):
            exc = Exception("XABLAU")
            handler = AsyncStreamHandler(
                level=10,
                stream=self.write_pipe,
                formatter=Mock(format=Mock(side_effect=exc)),
            )
            with asynctest.patch.object(
                handler, "handle_error"
            ) as handle_error:
                await handler.emit(self.record)

                handle_error.assert_awaited_once_with(self.record, exc)
                writer.write.assert_not_awaited()
                writer.drain.assert_not_awaited() 

def test_emit_writes_records_into_the_stream(self):
        msg = self.record.msg
        formatter = Mock(format=Mock(return_value=msg))
        writer = Mock(write=Mock(), drain=CoroutineMock())

        with patch(
            "aiologger.handlers.streams.StreamWriter", return_value=writer
        ):
            handler = AsyncStreamHandler(
                level=10, stream=self.write_pipe, formatter=formatter
            )

            await handler.emit(self.record)

            writer.write.assert_called_once_with(
                (msg + handler.terminator).encode()
            )
            writer.drain.assert_awaited_once()
            await handler.close() 

def test_watch_timeout_with_resource_version(self):
        fake_resp = CoroutineMock()
        fake_resp.content.readline = CoroutineMock()
        fake_resp.release = Mock()

        fake_resp.content.readline.side_effect = [asyncio.TimeoutError(),
                                                  b""]

        fake_api = Mock()
        fake_api.get_namespaces = CoroutineMock(return_value=fake_resp)
        fake_api.get_namespaces.__doc__ = ':return: V1NamespaceList'

        watch = kubernetes_asyncio.watch.Watch()
        async with watch.stream(fake_api.get_namespaces, resource_version='10') as stream:
            async for e in stream: # noqa
                pass

        # all calls use the passed resource version
        fake_api.get_namespaces.assert_has_calls(
            [call(_preload_content=False, watch=True, resource_version='10'),
             call(_preload_content=False, watch=True, resource_version='10')])

        fake_resp.release.assert_called_once_with()
        self.assertEqual(watch.resource_version, '10') 

def test_service_can_be_unregistered_by_id():
    with mock.patch('bookshelf.discolib.client') as client_func:
        client = asynctest.CoroutineMock()
        svc = mock.MagicMock()
        
        client_func.return_value = client
        client.agent.services.register.return_value = svc
        
        res = await discolib.register_service(id="svc1", name="svc",
                                              tags=["a", "b"],
                                              port=9000,
                                              address="localhost")
        assert res == svc

        client.agent.services.deregister.return_value = True
        res = await discolib.deregister_service("svc1")
        client.agent.services.deregister.assert_called_with("svc1")
        assert res is True 

def test_resources_utilization_nothing_special(
    mock_get_mesos_master, mock_get_resource_utilization_by_grouping
):
    request = testing.DummyRequest()
    request.swagger_data = {"groupings": None, "filter": None}
    mock_mesos_state = mock.Mock()
    mock_master = mock.Mock(
        state=asynctest.CoroutineMock(return_value=mock_mesos_state)
    )
    mock_get_mesos_master.return_value = mock_master

    mock_get_resource_utilization_by_grouping.return_value = {
        frozenset([("superregion", "unknown")]): {
            "total": metastatus_lib.ResourceInfo(cpus=10.0, mem=512.0, disk=100.0),
            "free": metastatus_lib.ResourceInfo(cpus=8.0, mem=312.0, disk=20.0),
        }
    }

    resp = resources_utilization(request)
    body = json.loads(resp.body.decode("utf-8"))

    assert resp.status_int == 200
    assert len(body) == 1
    assert set(body[0].keys()) == {"disk", "mem", "groupings", "cpus", "gpus"} 

def test_resources_utilization_with_grouping(mock_get_mesos_master):
    request = testing.DummyRequest()
    request.swagger_data = {"groupings": ["region", "pool"], "filter": None}
    mock_master = mock.Mock(
        state=asynctest.CoroutineMock(
            func=asynctest.CoroutineMock(),  # https://github.com/notion/a_sync/pull/40
            return_value=mock_mesos_state,
        )
    )
    mock_get_mesos_master.return_value = mock_master

    resp = resources_utilization(request)
    body = json.loads(resp.body.decode("utf-8"))

    assert resp.status_int == 200
    # 4 groupings, 2x2 attrs for 5 slaves
    assert len(body) == 4 

def test_get_cpu_usage_bad():
    fake_task = mock.create_autospec(mesos.task.Task)
    fake_task.cpu_limit = asynctest.CoroutineMock(return_value=1.1)
    fake_duration = 100
    fake_task.stats = asynctest.CoroutineMock(
        return_value={"cpus_system_time_secs": 50.0, "cpus_user_time_secs": 50.0}
    )
    current_time = datetime.datetime.now()
    fake_task.__getitem__.return_value = [
        {
            "state": "TASK_RUNNING",
            "timestamp": int(current_time.strftime("%s")) - fake_duration,
        }
    ]
    with asynctest.patch(
        "paasta_tools.mesos_tools.datetime.datetime", autospec=True
    ) as mock_datetime:
        mock_datetime.now.return_value = current_time
        actual = await mesos_tools.get_cpu_usage(fake_task)
    assert PaastaColors.red("100.0%") in actual 

def session(event_loop):
    fake_retry_policy = asynctest.MagicMock(wraps=aiozk.session.RetryPolicy.forever())
    session = aiozk.session.Session(
        'zookeeper.test',
        timeout=10,
        retry_policy=fake_retry_policy,
        allow_read_only=True,
        read_timeout=30,
        loop=asynctest.MagicMock(wraps=event_loop),
    )
    session.state.transition_to(aiozk.session.States.CONNECTED)
    session.conn = asynctest.MagicMock()
    session.conn.send = asynctest.CoroutineMock()
    session.conn.close = asynctest.CoroutineMock()
    session.ensure_safe_state = asynctest.CoroutineMock()
    session.set_heartbeat = mock.Mock()
    return session 

def test_is_draining_yes(self):
        fake_response = mock.Mock(
            status=503,
            text=asynctest.CoroutineMock(
                return_value="Service service in down state since 1435694078.778886 "
                "until 1435694178.780000: Drained by Paasta"
            ),
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response)
                )
            )
        ):
            assert await self.drain_method.is_draining(fake_task) is True 

def test_discovery_response_empty() -> None:
    """Test an empty discovery returns an empty list."""
    session = Mock()

    with patch("pydeconz.utils.async_request", new=CoroutineMock(return_value={})):
        response = await utils.async_discovery(session)

    assert not response 

def test_request() -> None:
    """Test a successful call of request."""
    response = Mock()
    response.content_type = "application/json"
    response.json = CoroutineMock(return_value={"json": "response"})
    session = CoroutineMock(return_value=response)

    result = await utils.async_request(session, "url")

    assert result == {"json": "response"} 

def test_all_packages(master: Master) -> None:
    expected = [["aiohttp", "", "", "", "69"]]
    master.rpc = asynctest.CoroutineMock(return_value=expected)  # type: ignore
    pacakges = await master.all_packages()
    assert expected == pacakges 

def test_request_fails_client_error() -> None:
    """Test a successful call of request."""
    session = CoroutineMock(side_effect=aiohttp.ClientError)

    with pytest.raises(errors.RequestError) as e_info:
        await utils.async_request(session, "url")

    assert str(e_info.value) == "Error requesting data from url: " 

def test_calls_backend_with_right_parameters_remove_user(self):
        account = Account(**ACCOUNT_DEV_DICT)
        user = User(**USER_WITH_NO_ACCOUNTS_DICT)
        backend = CoroutineMock(remove_user=CoroutineMock())
        await AccountsService.remove_user_from_account(user, account, backend)
        backend.remove_user.assert_awaited_with(user, account) 

def test_service_discovery_can_be_empty():
    with mock.patch('bookshelf.discolib.client') as client_func:
        client = asynctest.CoroutineMock()
        client_func.return_value = client
        client.catalog.services.return_value = {}
        res = await discolib.discover_services(["tag1"])
        assert len(res) == 0 

def test_discovery() -> None:
    """Test a successful call to discovery."""
    session = Mock()

    with patch(
        "pydeconz.utils.async_request",
        new=CoroutineMock(
            return_value=[
                {
                    "id": "123456FFFFABCDEF",
                    "internalipaddress": "host1",
                    "internalport": "port1",
                },
                {
                    "id": "234567BCDEFG",
                    "internalipaddress": "host2",
                    "internalport": "port2",
                },
            ]
        ),
    ):
        response = await utils.async_discovery(session)

    assert [
        {"bridgeid": "123456ABCDEF", "host": "host1", "port": "port1"},
        {"bridgeid": "234567BCDEFG", "host": "host2", "port": "port2"},
    ] == response 

def test_delete_all_keys() -> None:
    """Test a successful call of delete_all_keys.

    Delete all keys doesn't care what happens with delete_api_key.
    """
    session = Mock()

    with patch(
        "pydeconz.utils.async_request",
        new=CoroutineMock(return_value={"whitelist": {1: "123", 2: "456"}}),
    ):
        await utils.async_delete_all_keys(session, IP, PORT, API_KEY) 

def test_delete_api_key() -> None:
    """Test a successful call of delete_api_key."""
    session = Mock()

    with patch("pydeconz.utils.async_request", new=CoroutineMock(return_value=True)):
        await utils.async_delete_api_key(session, IP, PORT, API_KEY) 

def test_get_api_key_with_credentials() -> None:
    """Test a successful call of get_api_key with user crendentials."""
    session = Mock()

    with patch(
        "pydeconz.utils.async_request",
        new=CoroutineMock(return_value=[{"success": {"username": API_KEY}}]),
    ):
        response = await utils.async_get_api_key(
            session, IP, PORT, username="user", password="pass"
        )

    assert response == API_KEY 

def test_get_api_key() -> None:
    """Test a successful call of get_api_key."""
    session = Mock()

    with patch(
        "pydeconz.utils.async_request",
        new=CoroutineMock(return_value=[{"success": {"username": API_KEY}}]),
    ):
        response = await utils.async_get_api_key(session, IP, PORT)

    assert response == API_KEY 

def test_initialize(session) -> None:
    """Test a successful call of load_parameters."""
    with patch(
        "pydeconz.DeconzSession.request",
        new=CoroutineMock(
            return_value={
                "config": {"bridgeid": "012345"},
                "groups": {
                    "g1": {
                        "id": "gid",
                        "scenes": [{"id": "sc1", "name": "scene1"}],
                        "state": {},
                        "action": {},
                        "lights": [],
                    }
                },
                "lights": {"l1": {"state": {}}},
                "sensors": {
                    "s1": {"type": GenericStatus.ZHATYPE[0], "state": {}, "config": {}}
                },
            }
        ),
    ):
        await session.initialize()
    assert "g1" in session.groups
    assert session.groups["g1"].id == "gid"
    assert "sc1" in session.groups["g1"].scenes
    assert session.groups["g1"].deconz_id == "/groups/g1"
    assert session.groups["g1"].scenes["sc1"].id == "sc1"
    assert "l1" in session.lights
    assert session.lights["l1"].deconz_id == "/lights/l1"
    assert "s1" in session.sensors
    assert session.sensors["s1"].deconz_id == "/sensors/s1"
    assert session.sensors["s1"].type == GenericStatus.ZHATYPE[0] 

def test_get_json_body_from_service():
    session = asynctest.MagicMock()
    session.get.return_value.__aenter__.return_value.json = asynctest.CoroutineMock(
        return_value=mock.sentinel.json_body
    )
    assert (
        await autoscaling_service_lib.get_json_body_from_service(
            "fake-host", "fake-port", "fake-endpoint", session
        )
        == mock.sentinel.json_body
    )
    session.get.assert_called_once_with(
        "http://fake-host:fake-port/fake-endpoint", headers={"User-Agent": mock.ANY}
    ) 

def test_is_draining_no(self):
        fake_response = mock.Mock(
            status=200, text=asynctest.CoroutineMock(return_value="")
        )
        fake_task = mock.Mock(host="fake_host", ports=[54321])
        with mock_ClientSession(
            get=mock.Mock(
                return_value=asynctest.MagicMock(
                    __aenter__=asynctest.CoroutineMock(return_value=fake_response)
                )
            )
        ):
            assert await self.drain_method.is_draining(fake_task) is False 

def test_get_files_for_tasks_some():
    mock_task = asynctest.MagicMock(spec=task.Task)
    mock_file = Mock()
    mock_file_2 = Mock()
    mock_file.exists = asynctest.CoroutineMock(return_value=False)
    mock_file_2.exists = asynctest.CoroutineMock(return_value=True)
    mock_task.file.side_effect = [mock_file, mock_file_2]
    files = cluster.get_files_for_tasks([mock_task], ["myfile", "myotherfile"], 1)
    files = a_sync.block(aiter_to_list, files)
    assert files == [mock_file_2]