OwlCyberSecurity - MANAGER

Edit File: test_contextvars.py

from __future__ import print_function

import gc
import sys
import unittest

from functools import partial
from unittest import skipUnless
from unittest import skipIf

from greenlet import greenlet
from greenlet import getcurrent
from . import TestCase

try:
    from contextvars import Context
    from contextvars import ContextVar
    from contextvars import copy_context
    # From the documentation:
    #
    # Important: Context Variables should be created at the top module
    # level and never in closures. Context objects hold strong
    # references to context variables which prevents context variables
    # from being properly garbage collected.
    ID_VAR = ContextVar("id", default=None)
    VAR_VAR = ContextVar("var", default=None)
    ContextVar = None
except ImportError:
    Context = ContextVar = copy_context = None

# We don't support testing if greenlet's built-in context var support is disabled.
@skipUnless(Context is not None, "ContextVar not supported")
class ContextVarsTests(TestCase):
    def _new_ctx_run(self, *args, **kwargs):
        return copy_context().run(*args, **kwargs)

def _increment(self, greenlet_id, callback, counts, expect):
        ctx_var = ID_VAR
        if expect is None:
            self.assertIsNone(ctx_var.get())
        else:
            self.assertEqual(ctx_var.get(), expect)
        ctx_var.set(greenlet_id)
        for _ in range(2):
            counts[ctx_var.get()] += 1
            callback()

def _test_context(self, propagate_by):
        # pylint:disable=too-many-branches
        ID_VAR.set(0)

callback = getcurrent().switch
        counts = dict((i, 0) for i in range(5))

lets = [
            greenlet(partial(
                partial(
                    copy_context().run,
                    self._increment
                ) if propagate_by == "run" else self._increment,
                greenlet_id=i,
                callback=callback,
                counts=counts,
                expect=(
                    i - 1 if propagate_by == "share" else
                    0 if propagate_by in ("set", "run") else None
                )
            ))
            for i in range(1, 5)
        ]

for let in lets:
            if propagate_by == "set":
                let.gr_context = copy_context()
            elif propagate_by == "share":
                let.gr_context = getcurrent().gr_context

for i in range(2):
            counts[ID_VAR.get()] += 1
            for let in lets:
                let.switch()

if propagate_by == "run":
            # Must leave each context.run() in reverse order of entry
            for let in reversed(lets):
                let.switch()
        else:
            # No context.run(), so fine to exit in any order.
            for let in lets:
                let.switch()

for let in lets:
            self.assertTrue(let.dead)
            # When using run(), we leave the run() as the greenlet dies,
            # and there's no context "underneath". When not using run(),
            # gr_context still reflects the context the greenlet was
            # running in.
            if propagate_by == 'run':
                self.assertIsNone(let.gr_context)
            else:
                self.assertIsNotNone(let.gr_context)

if propagate_by == "share":
            self.assertEqual(counts, {0: 1, 1: 1, 2: 1, 3: 1, 4: 6})
        else:
            self.assertEqual(set(counts.values()), set([2]))

def test_context_propagated_by_context_run(self):
        self._new_ctx_run(self._test_context, "run")

def test_context_propagated_by_setting_attribute(self):
        self._new_ctx_run(self._test_context, "set")

def test_context_not_propagated(self):
        self._new_ctx_run(self._test_context, None)

def test_context_shared(self):
        self._new_ctx_run(self._test_context, "share")

def test_break_ctxvars(self):
        let1 = greenlet(copy_context().run)
        let2 = greenlet(copy_context().run)
        let1.switch(getcurrent().switch)
        let2.switch(getcurrent().switch)
        # Since let2 entered the current context and let1 exits its own, the
        # interpreter emits:
        # RuntimeError: cannot exit context: thread state references a different context object
        let1.switch()

def test_not_broken_if_using_attribute_instead_of_context_run(self):
        let1 = greenlet(getcurrent().switch)
        let2 = greenlet(getcurrent().switch)
        let1.gr_context = copy_context()
        let2.gr_context = copy_context()
        let1.switch()
        let2.switch()
        let1.switch()
        let2.switch()

def test_context_assignment_while_running(self):
        # pylint:disable=too-many-statements
        ID_VAR.set(None)

def target():
            self.assertIsNone(ID_VAR.get())
            self.assertIsNone(gr.gr_context)

# Context is created on first use
            ID_VAR.set(1)
            self.assertIsInstance(gr.gr_context, Context)
            self.assertEqual(ID_VAR.get(), 1)
            self.assertEqual(gr.gr_context[ID_VAR], 1)

# Clearing the context makes it get re-created as another
            # empty context when next used
            old_context = gr.gr_context
            gr.gr_context = None  # assign None while running
            self.assertIsNone(ID_VAR.get())
            self.assertIsNone(gr.gr_context)
            ID_VAR.set(2)
            self.assertIsInstance(gr.gr_context, Context)
            self.assertEqual(ID_VAR.get(), 2)
            self.assertEqual(gr.gr_context[ID_VAR], 2)

new_context = gr.gr_context
            getcurrent().parent.switch((old_context, new_context))
            # parent switches us back to old_context

self.assertEqual(ID_VAR.get(), 1)
            gr.gr_context = new_context  # assign non-None while running
            self.assertEqual(ID_VAR.get(), 2)

getcurrent().parent.switch()
            # parent switches us back to no context
            self.assertIsNone(ID_VAR.get())
            self.assertIsNone(gr.gr_context)
            gr.gr_context = old_context
            self.assertEqual(ID_VAR.get(), 1)

getcurrent().parent.switch()
            # parent switches us back to no context
            self.assertIsNone(ID_VAR.get())
            self.assertIsNone(gr.gr_context)

gr = greenlet(target)

with self.assertRaisesRegex(AttributeError, "can't delete context attribute"):
            del gr.gr_context

self.assertIsNone(gr.gr_context)
        old_context, new_context = gr.switch()
        self.assertIs(new_context, gr.gr_context)
        self.assertEqual(old_context[ID_VAR], 1)
        self.assertEqual(new_context[ID_VAR], 2)
        self.assertEqual(new_context.run(ID_VAR.get), 2)
        gr.gr_context = old_context  # assign non-None while suspended
        gr.switch()
        self.assertIs(gr.gr_context, new_context)
        gr.gr_context = None  # assign None while suspended
        gr.switch()
        self.assertIs(gr.gr_context, old_context)
        gr.gr_context = None
        gr.switch()
        self.assertIsNone(gr.gr_context)

# Make sure there are no reference leaks
        gr = None
        gc.collect()
        self.assertEqual(sys.getrefcount(old_context), 2)
        self.assertEqual(sys.getrefcount(new_context), 2)

def test_context_assignment_different_thread(self):
        import threading
        VAR_VAR.set(None)
        ctx = Context()

is_running = threading.Event()
        should_suspend = threading.Event()
        did_suspend = threading.Event()
        should_exit = threading.Event()
        holder = []

def greenlet_in_thread_fn():
            VAR_VAR.set(1)
            is_running.set()
            should_suspend.wait(10)
            VAR_VAR.set(2)
            getcurrent().parent.switch()
            holder.append(VAR_VAR.get())

def thread_fn():
            gr = greenlet(greenlet_in_thread_fn)
            gr.gr_context = ctx
            holder.append(gr)
            gr.switch()
            did_suspend.set()
            should_exit.wait(10)
            gr.switch()
            del gr
            greenlet() # trigger cleanup

thread = threading.Thread(target=thread_fn, daemon=True)
        thread.start()
        is_running.wait(10)
        gr = holder[0]

# Can't access or modify context if the greenlet is running
        # in a different thread
        with self.assertRaisesRegex(ValueError, "running in a different"):
            getattr(gr, 'gr_context')
        with self.assertRaisesRegex(ValueError, "running in a different"):
            gr.gr_context = None

should_suspend.set()
        did_suspend.wait(10)

# OK to access and modify context if greenlet is suspended
        self.assertIs(gr.gr_context, ctx)
        self.assertEqual(gr.gr_context[VAR_VAR], 2)
        gr.gr_context = None

should_exit.set()
        thread.join(10)

self.assertEqual(holder, [gr, None])

# Context can still be accessed/modified when greenlet is dead:
        self.assertIsNone(gr.gr_context)
        gr.gr_context = ctx
        self.assertIs(gr.gr_context, ctx)

# Otherwise we leak greenlets on some platforms.
        # XXX: Should be able to do this automatically
        del holder[:]
        gr = None
        thread = None

def test_context_assignment_wrong_type(self):
        g = greenlet()
        with self.assertRaisesRegex(TypeError,
                                    "greenlet context must be a contextvars.Context or None"):
            g.gr_context = self

@skipIf(Context is not None, "ContextVar supported")
class NoContextVarsTests(TestCase):
    def test_contextvars_errors(self):
        let1 = greenlet(getcurrent().switch)
        self.assertFalse(hasattr(let1, 'gr_context'))
        with self.assertRaises(AttributeError):
            getattr(let1, 'gr_context')

with self.assertRaises(AttributeError):
            let1.gr_context = None

let1.switch()

with self.assertRaises(AttributeError):
            getattr(let1, 'gr_context')

with self.assertRaises(AttributeError):
            let1.gr_context = None

del let1

if __name__ == '__main__':
    unittest.main()