Python MagicMock.cursor方法代码示例

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_grant_user_roles(self, pgidentifier, ensure_role, log):
        pgidentifier.side_effect = lambda d: "q_{}".format(d)

        existing_roles = set(["roleA", "roleB"])
        wanted_roles = set(["roleB", "roleC"])

        con = MagicMock()
        cur = con.cursor()
        cur.fetchall.return_value = [(r,) for r in existing_roles]

        postgresql.grant_user_roles(con, "fred", wanted_roles)

        # A new role was ensured. The others we know exist.
        ensure_role.assert_called_once_with(con, "roleC")

        role_query = dedent(
            """\
            SELECT role.rolname
            FROM
                pg_roles AS role,
                pg_roles AS member,
                pg_auth_members
            WHERE
                member.oid = pg_auth_members.member
                AND role.oid = pg_auth_members.roleid
                AND member.rolname = %s
            """
        )
        cur.execute.assert_has_calls([call(role_query, ("fred",)), call("GRANT %s TO %s", ("q_roleC", "q_fred"))])

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_ensure_user(self, role_exists, pgidentifier):
        con = MagicMock()
        cur = con.cursor()

        # Create a new boring user
        role_exists.return_value = False
        pgidentifier.return_value = sentinel.quoted_user
        postgresql.ensure_user(con, sentinel.user, sentinel.secret)
        pgidentifier.assert_called_once_with(sentinel.user)
        cur.execute.assert_called_once_with(
            "CREATE ROLE %s WITH LOGIN NOSUPERUSER NOREPLICATION PASSWORD %s", (sentinel.quoted_user, sentinel.secret)
        )

        # Ensure an existing user is a superuser
        role_exists.return_value = True
        cur.execute.reset_mock()
        postgresql.ensure_user(con, sentinel.user, sentinel.secret, superuser=True)
        cur.execute.assert_called_once_with(
            "ALTER ROLE %s WITH LOGIN SUPERUSER NOREPLICATION PASSWORD %s", (sentinel.quoted_user, sentinel.secret)
        )

        # Create a new user with replication permissions.
        role_exists.return_value = False
        cur.execute.reset_mock()
        postgresql.ensure_user(con, sentinel.user, sentinel.secret, replication=True)
        cur.execute.assert_called_once_with(
            "CREATE ROLE %s WITH LOGIN NOSUPERUSER REPLICATION PASSWORD %s", (sentinel.quoted_user, sentinel.secret)
        )

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_ensure_extensions(self, pgidentifier, log):
        con = MagicMock()
        cur = con.cursor()

        pgidentifier.side_effect = lambda d: "q_{}".format(d)

        existing_extensions = set(["extA", "extB"])
        wanted_extensions = set(["extB", "extC"])

        cur.fetchall.return_value = [[x] for x in existing_extensions]
        postgresql.ensure_extensions(con, wanted_extensions)
        cur.execute.assert_has_calls(
            [call("SELECT extname FROM pg_extension"), call("CREATE EXTENSION %s", ("q_extC",))]
        )

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_role_exists(self):
        con = MagicMock()
        cur = con.cursor()

        # Exists
        cur.fetchone.return_value = sentinel.something
        self.assertTrue(postgresql.role_exists(con, sentinel.role))
        cur.execute.assert_called_once_with("SELECT TRUE FROM pg_roles WHERE rolname=%s", (sentinel.role,))

        # Does not exist
        cur.fetchone.return_value = None
        cur.execute.reset_mock()
        self.assertFalse(postgresql.role_exists(con, sentinel.role))
        cur.execute.assert_called_once_with("SELECT TRUE FROM pg_roles WHERE rolname=%s", (sentinel.role,))

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_ensure_role(self, pgidentifier):
        con = MagicMock()
        cur = con.cursor()

        pgidentifier.side_effect = lambda d: "q_{}".format(d)

        # If the role already exists, nothing happens.
        cur.fetchone.return_value = sentinel.something
        postgresql.ensure_role(con, "roleA")
        cur.execute.assert_called_once_with("SELECT TRUE FROM pg_roles WHERE rolname=%s", ("roleA",))

        # If the role does not exist, it is created.
        cur.fetchone.return_value = None
        postgresql.ensure_role(con, "roleA")
        cur.execute.assert_has_calls([call("CREATE ROLE %s INHERIT NOLOGIN", ("q_roleA",))])

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_grant_database_privileges(self):
        con = MagicMock()
        cur = con.cursor()
        privs = ["privA", "privB"]
        postgresql.grant_database_privileges(con, "a_Role", "a_DB", privs)

        cur.execute.assert_has_calls(
            [
                call(
                    "GRANT %s ON DATABASE %s TO %s",
                    (
                        postgresql.AsIs("privA"),  # Unquoted. Its a keyword.
                        postgresql.AsIs('"a_DB"'),
                        postgresql.AsIs('"a_Role"'),
                    ),
                ),
                call(
                    "GRANT %s ON DATABASE %s TO %s",
                    (postgresql.AsIs("privB"), postgresql.AsIs('"a_DB"'), postgresql.AsIs('"a_Role"')),
                ),
            ]
        )

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
    def test_execute(self, mock_data_lake_hook, mock_oracle_hook):
        task_id = "some_test_id"
        sql = "some_sql"
        sql_params = {':p_data': "2018-01-01"}
        oracle_conn_id = "oracle_conn_id"
        filename = "some_filename"
        azure_data_lake_conn_id = 'azure_data_lake_conn_id'
        azure_data_lake_path = 'azure_data_lake_path'
        delimiter = '|'
        encoding = 'latin-1'
        cursor_description = [
            ('id', "<class 'cx_Oracle.NUMBER'>", 39, None, 38, 0, 0),
            ('description', "<class 'cx_Oracle.STRING'>", 60, 240, None, None, 1)
        ]
        cursor_rows = [[1, 'description 1'], [2, 'description 2']]
        cursor_mock = MagicMock()
        cursor_mock.description.return_value = cursor_description
        cursor_mock.__iter__.return_value = cursor_rows
        mock_oracle_conn = MagicMock()
        mock_oracle_conn.cursor().return_value = cursor_mock
        mock_oracle_hook.get_conn().return_value = mock_oracle_conn

        op = OracleToAzureDataLakeTransfer(
            task_id=task_id,
            filename=filename,
            oracle_conn_id=oracle_conn_id,
            sql=sql,
            sql_params=sql_params,
            azure_data_lake_conn_id=azure_data_lake_conn_id,
            azure_data_lake_path=azure_data_lake_path,
            delimiter=delimiter,
            encoding=encoding)

        op.execute(None)

        mock_oracle_hook.assert_called_once_with(oracle_conn_id=oracle_conn_id)
        mock_data_lake_hook.assert_called_once_with(
            azure_data_lake_conn_id=azure_data_lake_conn_id)

# 需要导入模块: from unittest.mock import MagicMock [as 别名]
# 或者: from unittest.mock.MagicMock import cursor [as 别名]
def get_mock_connection():
    mock_connection = MagicMock(spec=connection)
    mock_connection.cursor = lambda: MagicMock(spec=cursor)
    return mock_connection