Example #1
0
    def test_remove_many_invalid_keys(self, m_gcp):
        """
        Test failure when one or more provided keys are not present in
        store.
        """
        s = PostgresKeyValueStore(create_table=False)

        # Simulate the batch execute returning nothing.  This simulates no
        # rows being found by the first call to the method when checking
        # for key presence in table.
        s._check_contained_keys = mock.Mock(return_value={0, 1})
        PY2_SET_KEY_ERROR_RE = r"set\(\[(?:0|1), (?:0|1)\]\)"
        PY3_SET_KEY_ERROR_RE = r"{(?:0|1), (?:0|1)}"
        SET_KEY_ERROR_RE = r'^(?:{}|{})$'.format(PY2_SET_KEY_ERROR_RE,
                                                 PY3_SET_KEY_ERROR_RE)
        with pytest.raises(KeyError, match=SET_KEY_ERROR_RE):
            s.remove_many([0, 1])

        # Simulate only one of the keys existing in the table.
        s._check_contained_keys = mock.Mock(return_value={1})
        with pytest.raises(KeyError, match=r'^1$'):
            s.remove_many([0, 1])
        s._check_contained_keys = mock.Mock(return_value={0})
        with pytest.raises(KeyError, match=r'^0$'):
            s.remove_many([0, 1])
        def test_remove_many_invalid_keys(self, m_gcp):
            """
            Test failure when one or more provided keys are not present in
            store.
            """
            s = PostgresKeyValueStore(create_table=False)

            # Simulate the batch execute returning nothing.  This simulates no
            # rows being found by the first call to the method when checking
            # for key presence in table.
            s._check_contained_keys = mock.Mock(return_value={0, 1})
            PY2_SET_KEY_ERROR_RE = "set\(\[(?:0|1), (?:0|1)\]\)"
            PY3_SET_KEY_ERROR_RE = "{(?:0|1), (?:0|1)}"
            self.assertRaisesRegexp(
                KeyError, '^(?:{}|{})$'.format(PY2_SET_KEY_ERROR_RE,
                                               PY3_SET_KEY_ERROR_RE),
                s.remove_many, [0, 1]
            )

            # Simulate only one of the keys existing in the table.
            s._check_contained_keys = mock.Mock(return_value={1})
            self.assertRaisesRegexp(
                KeyError, '^1$',
                s.remove_many, [0, 1]
            )
            s._check_contained_keys = mock.Mock(return_value={0})
            self.assertRaisesRegexp(
                KeyError, '^0$',
                s.remove_many, [0, 1]
            )
        def test_remove_many(self, m_psqlExecBatch, m_gcp):
            """
            Test expected calls to psql `execute_batch` function when removing
            multiple items.
            """
            expected_key_1 = 'test_remove_many_key_1'
            exp_key_1_bytea = PostgresKeyValueStore._py_to_bin(expected_key_1)
            expected_key_2 = 'test_remove_many_key_2'
            exp_key_2_bytea = PostgresKeyValueStore._py_to_bin(expected_key_2)

            # Skip table creation calls for simplicity.
            s = PostgresKeyValueStore(create_table=False)

            # Mock PSQL cursor stuff because we aren't actually connecting to a
            # database.
            # - `get_psql_connection` uses `smqtk.utils.postgres.
            #   get_connection_pool`, so the return is a mock object.
            # - Cursor is created via a context (i.e. __enter__()) when utilized
            #   in `PsqlConnectionHelper` execute methods.
            #: :type: mock.Mock
            mock_cursor = s._psql_helper.get_psql_connection().cursor() \
                .__enter__()

            # Mocking `PostgresKeyValueStore` key-check method so as to pretend
            # that the given keys exist in the database
            s._check_contained_keys = mock.MagicMock(return_value=set())

            s.remove_many([expected_key_1, expected_key_2])

            # As a result of this call, we expect:
            # - ``psycopg2.extras.execute_batch`` should have been called once
            #   when deleting key-value pairs in db (2 < `s._batch_size`)
            #
            # We back to break up the argument equality check of recorded mock
            # function call arguments due to ``psycopg2.Binary`` instances not
            # being comparable.

            m_psqlExecBatch.assert_called_once()

            # Confirm call arguments provided to
            # ``psycopg2.extras.execute_batch`` are as expected.
            # -
            expected_del_q = "DELETE FROM data_set WHERE key LIKE %(key_like)s"
            psqlExecBatch_call_args = m_psqlExecBatch.call_args[0]
            psqlExecBatch_kwargs = m_psqlExecBatch.call_args[1]
            self.assertEqual(psqlExecBatch_call_args[0], mock_cursor)
            self.assertEqual(psqlExecBatch_call_args[1], expected_del_q)
            # 3rd argument is a list of dictionaries for 'key_like' replacements
            # - dictionary values are `psycopg2.extensions.Binary` type, which
            #   are not directly comparable (different instances). Have to
            #   convert to bytes representation in order to compare.
            self.assertEqual(len(psqlExecBatch_call_args[2]), 2)
            self.assertSetEqual(
                {
                    d['key_like'].getquoted()
                    for d in psqlExecBatch_call_args[2]
                }, {exp_key_1_bytea.getquoted(),
                    exp_key_2_bytea.getquoted()})
            self.assertIn('page_size', psqlExecBatch_kwargs)
            self.assertEqual(psqlExecBatch_kwargs['page_size'], s._batch_size)
            self.assertEqual(psqlExecBatch_kwargs['page_size'],
                             1000)  # default