google-cloud-testutils 1.6.1__py3-none-any.whl

test_utils/orchestrate.py

@@ -0,0 +1,446 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import math
+import queue
+import sys
+import threading
+import tokenize
+
+
+def orchestrate(*tests, **kwargs):
+    """
+    Orchestrate a deterministic concurrency test.
+
+    Runs test functions in separate threads, with each thread taking turns running up
+    until predefined syncpoints in a deterministic order. All possible orderings are
+    tested.
+
+    Most of the time, we try to use logic, best practices, and static analysis to insure
+    correct operation of concurrent code. Sometimes our powers of reasoning fail us and,
+    either through non-determistic stress testing or running code in production, a
+    concurrent bug is discovered. When this occurs, we'd like to have a regression test
+    to insure we've understood the problem and implemented a correct solution.
+    `orchestrate` provides a means of deterministically testing concurrent code so we
+    can write robust regression tests for complex concurrent scenarios.
+
+    `orchestrate` runs each passed in test function in its own thread. Threads then
+    "take turns" running. Turns are defined by setting syncpoints in the code under
+    test, using comments containing "pragma: SYNCPOINT". `orchestrate` will scan the
+    code under test and add syncpoints where it finds these comments.
+
+    For example, let's say you have the following code in production::
+
+        def hither_and_yon(destination):
+            hither(destination)
+            yon(destination)
+
+    You've found there's a concurrency bug when two threads execute this code with the
+    same argument, and you think that by adding a syncpoint between the calls to
+    `hither` and `yon` you can reproduce the problem in a regression test. First add a
+    comment with "pragma: SYNCPOINT" to the code under test::
+
+        def hither_and_yon(destination):
+            hither(destination)   # pragma: SYNCPOINT
+            yon(destination)
+
+    When testing with orchestrate, there will now be a syncpoint, or a pause, after the
+    call to `hither` and before the call to `yon`. Now you can write a test to exercise
+    `hither_and_yon` running in parallel::
+
+        from unittest import mock
+        from tests.unit import orchestrate
+
+        from myorg.myproj.sales import travel
+
+        def test_concurrent_hither_and_yon():
+
+            def test_hither_and_yon():
+                assert something
+                travel.hither_and_yon("Raleigh")
+                assert something_else
+
+            counts = orchestrate.orchestrate(test_hither_and_yon, test_hither_and_yon)
+            assert counts == (2, 2)
+
+    What `orchestrate` will do now is take each of the two test functions passed in
+    (actually the same function, twice, in this case), run them serially, and count the
+    number of turns it takes to run each test to completion. In this example, it will
+    take two turns for each test: one turn to start the thread and execute up until the
+    syncpoint, and then another turn to execute from the syncpoint to the end of the
+    test. The number of turns will always be one greater than the number of syncpoints
+    encountered when executing the test.
+
+    Once the counts have been taken, `orchestrate` will construct a test sequence that
+    represents all of the turns taken by the passed in tests, with each value in the
+    sequence representing the index of the test whose turn it is in the sequence. In
+    this example, then, it would produce::
+
+        [0, 0, 1, 1]
+
+    This represents the first test taking both of its turns, followed by the second test
+    taking both of its turns. At this point this scenario has already been tested,
+    because this is what was run to produce the counts and the initial test sequence.
+    Now `orchestrate` will run all of the remaining scenarios by finding all the
+    permutations of the test sequence and executing those, in turn::
+
+        [0, 1, 0, 1]
+        [0, 1, 1, 0]
+        [1, 0, 0, 1]
+        [1, 0, 1, 0]
+        [1, 1, 0, 0]
+
+    You'll notice in our example that since both test functions are actually the same
+    function, that although it tested 6 scenarios there are effectively only really 3
+    unique scenarios. For the time being, though, `orchestrate` doesn't attempt to
+    detect this condition or optimize for it.
+
+    There are some performance considerations that should be taken into account when
+    writing tests. The number of unique test sequences grows quite quickly with the
+    number of turns taken by the functions under test. Our simple example with two
+    threads each taking two turns, only yielded 6 scenarios, but two threads each taking
+    6 turns, for example, yields 924 scenarios. Add another six step thread, for a total
+    of three threads, and now you have over 17 thousand scenarios. In general, use the
+    least number of steps/threads you can get away with and still expose the behavior
+    you want to correct.
+
+    For the same reason as above, it is recommended that if you have many concurrent
+    tests, that you name your syncpoints so that you're not accidentally using
+    syncpoints intended for other tests, as this will add steps to your tests. While
+    it's not problematic from a testing standpoint to have extra steps in your tests, it
+    can use computing resources unnecessarily. A name can be added to any syncpoint
+    after the `SYNCPOINT` keyword in the pragma definition::
+
+        def hither_and_yon(destination):
+            hither(destination)   # pragma: SYNCPOINT hither and yon
+            yon(destination)
+
+    In your test, then, pass that name to `orchestrate` to cause it to use only
+    syncpoints with that name::
+
+        orchestrate.orchestrate(
+            test_hither_and_yon, test_hither_and_yon, name="hither and yon"
+        )
+
+    As soon as any error or failure is detected, no more scenarios are run
+    and that error is propagated to the main thread.
+
+    One limitation of `orchestrate` is that it cannot really be used with `coverage`,
+    since both tools use `sys.set_trace`. Any code that needs verifiable test coverage
+    should have additional tests that do not use `orchestrate`, since code that is run
+    under orchestrate will not show up in a coverage report generated by `coverage`.
+
+    Args:
+        tests (Tuple[Callable]): Test functions to be run. These functions will not be
+            called with any arguments, so they must not have any required arguments.
+        name (Optional[str]): Only use syncpoints with the given name. If omitted, only
+            unnamed syncpoints will be used.
+
+    Returns:
+        Tuple[int]: A tuple of the count of the number turns for test passed in. Can be
+            used a sanity check in tests to make sure you understand what's actually
+            happening during a test.
+    """
+    name = kwargs.pop("name", None)
+    if kwargs:
+        raise TypeError(
+            "Unexpected keyword arguments: {}".format(", ".join(kwargs.keys()))
+        )
+
+    # Produce an initial test sequence. The fundamental question we're always trying to
+    # answer is "whose turn is it?" First we'll find out how many "turns" each test
+    # needs to complete when run serially and use that to construct a sequence of
+    # indexes. When a test's index appears in the sequence, it is that test's turn to
+    # run. We'll start by constructing a sequence that would run each test through to
+    # completion serially, one after the other.
+    test_sequence = []
+    counts = []
+    for index, test in enumerate(tests):
+        thread = _TestThread(test, name)
+        for count in itertools.count(1):  # pragma: NO BRANCH
+            # Pragma is required because loop never finishes naturally.
+            thread.go()
+            if thread.finished:
+                break
+
+        counts.append(count)
+        test_sequence += [index] * count
+
+    # Now we can take that initial sequence and generate all of its permutations,
+    # running each one to try to uncover concurrency bugs
+    sequences = iter(_permutations(test_sequence))
+
+    # We already tested the first sequence getting our counts, so we can discard it
+    next(sequences)
+
+    # Test each sequence
+    for test_sequence in sequences:
+        threads = [_TestThread(test, name) for test in tests]
+        try:
+            for index in test_sequence:
+                threads[index].go()
+
+            # Its possible for number of turns to vary from one test run to the other,
+            # especially if there is some undiscovered concurrency bug. Go ahead and
+            # finish running each test to completion, if not already complete.
+            for thread in threads:
+                while not thread.finished:
+                    thread.go()
+
+        except Exception:
+            # If an exception occurs, we still need to let any threads that are still
+            # going finish up. Additional exceptions are silently ignored.
+            for thread in threads:
+                thread.finish()
+            raise
+
+    return tuple(counts)
+
+
+_local = threading.local()
+
+
+class _Conductor:
+    """Coordinate communication between main thread and a test thread.
+
+    Two way communicaton is maintained between the main thread and a test thread using
+    two synchronized queues (`queue.Queue`) each with a size of one.
+    """
+
+    def __init__(self):
+        self._notify = queue.Queue(1)
+        self._go = queue.Queue(1)
+
+    def notify(self):
+        """Called from test thread to let us know it's finished or is ready for its next
+        turn."""
+        self._notify.put(None)
+
+    def standby(self):
+        """Called from test thread in order to block until told to go."""
+        self._go.get()
+
+    def wait(self):
+        """Called from main thread to wait for test thread to either get to the
+        next syncpoint or finish."""
+        self._notify.get()
+
+    def go(self):
+        """Called from main thread to tell test thread to go."""
+        self._go.put(None)
+
+
+_SYNCPOINTS = {}
+"""Dict[str, Dict[str, Set[int]]]: Dict mapping source fileneme to a dict mapping
+syncpoint name to set of line numbers where syncpoints with that name occur in the
+source file.
+"""
+
+
+def _get_syncpoints(filename):
+    """Find syncpoints in a source file.
+
+    Does a simple tokenization of the source file, looking for comments with "pragma:
+    SYNCPOINT", and populates _SYNCPOINTS using the syncpoint name and line number in
+    the source file.
+    """
+    _SYNCPOINTS[filename] = syncpoints = {}
+
+    # Use tokenize to find pragma comments
+    with open(filename, "r") as pyfile:
+        tokens = tokenize.generate_tokens(pyfile.readline)
+        for type, value, start, end, line in tokens:
+            if type == tokenize.COMMENT and "pragma: SYNCPOINT" in value:
+                name = value.split("SYNCPOINT", 1)[1].strip()
+                if not name:
+                    name = None
+
+                if name not in syncpoints:
+                    syncpoints[name] = set()
+
+                lineno, column = start
+                syncpoints[name].add(lineno)
+
+
+class _TestThread:
+    """A thread for a test function."""
+
+    thread = None
+    finished = False
+    error = None
+    at_syncpoint = False
+
+    def __init__(self, test, name):
+        self.test = test
+        self.name = name
+        self.conductor = _Conductor()
+
+    def _run(self):
+        sys.settrace(self._trace)
+        _local.conductor = self.conductor
+        try:
+            self.test()
+        except Exception as error:
+            self.error = error
+        finally:
+            self.finished = True
+            self.conductor.notify()
+
+    def _sync(self):
+        # Tell main thread we're finished, for now
+        self.conductor.notify()
+
+        # Wait for the main thread to tell us to go again
+        self.conductor.standby()
+
+    def _trace(self, frame, event, arg):
+        """Argument to `sys.settrace`.
+
+        Handles frames during test run, syncing at syncpoints, when found.
+
+        Returns:
+            `None` if no more tracing is required for the function call, `self._trace`
+            if tracing should continue.
+        """
+        if self.at_syncpoint:
+            # We hit a syncpoint on the previous call, so now we sync.
+            self._sync()
+            self.at_syncpoint = False
+
+        filename = frame.f_globals.get("__file__")
+        if not filename:
+            # Can't trace code without a source file
+            return
+
+        if filename.endswith(".pyc"):
+            filename = filename[:-1]
+
+        if filename not in _SYNCPOINTS:
+            _get_syncpoints(filename)
+
+        syncpoints = _SYNCPOINTS[filename].get(self.name)
+        if not syncpoints:
+            # This file doesn't contain syncpoints, don't continue to trace
+            return
+
+        # We've hit a syncpoint. Execute whatever line the syncpoint is on and then
+        # sync next time this gets called.
+        if frame.f_lineno in syncpoints:
+            self.at_syncpoint = True
+
+        return self._trace
+
+    def go(self):
+        if self.finished:
+            return
+
+        if self.thread is None:
+            self.thread = threading.Thread(target=self._run)
+            self.thread.start()
+
+        else:
+            self.conductor.go()
+
+        self.conductor.wait()
+
+        if self.error:
+            raise self.error
+
+    def finish(self):
+        while not self.finished:
+            try:
+                self.go()
+            except Exception:
+                pass
+
+
+class _permutations:
+    """Generates a sequence of all permutations of `sequence`.
+
+    Permutations are returned in lexicographic order using the "Generation in
+    lexicographic order" algorithm described in `the Wikipedia article on "Permutation"
+    <https://en.wikipedia.org/wiki/Permutation>`_.
+
+    This implementation differs significantly from `itertools.permutations` in that the
+    value of individual elements is taken into account, thus eliminating redundant
+    orderings that would be produced by `itertools.permutations`.
+
+    Args:
+        sequence (Sequence[Any]): Sequence must be finite and orderable.
+
+    Returns:
+        Sequence[Sequence[Any]]: Set of all permutations of `sequence`.
+    """
+
+    def __init__(self, sequence):
+        self._start = tuple(sorted(sequence))
+
+    def __len__(self):
+        """Compute the number of permutations.
+
+        Let the number of elements in a sequence N and the number of repetitions for
+        individual members of the sequence be n1, n2, ... nx. The number of unique
+        permutations is: N! / n1! / n2! / ... / nx!.
+
+        For example, let `sequence` be [1, 2, 3, 1, 2, 3, 1, 2, 3]. The number of unique
+        permutations is: 9! / 3! / 3! / 3! = 1680.
+
+        See: "Permutations of multisets" in `the Wikipedia article on "Permutation"
+        <https://en.wikipedia.org/wiki/Permutation>`_.
+        """
+        repeats = [len(list(group)) for value, group in itertools.groupby(self._start)]
+        length = math.factorial(len(self._start))
+        for repeat in repeats:
+            length /= math.factorial(repeat)
+
+        return int(length)
+
+    def __iter__(self):
+        """Iterate over permutations.
+
+        See: "Generation in lexicographic order" algorithm described in `the Wikipedia
+        article on "Permutation" <https://en.wikipedia.org/wiki/Permutation>`_.
+        """
+        current = list(self._start)
+        size = len(current)
+
+        while True:
+            yield tuple(current)
+
+            # 1. Find the largest index i such that a[i] < a[i + 1].
+            for i in range(size - 2, -1, -1):
+                if current[i] < current[i + 1]:
+                    break
+
+            else:
+                # If no such index exists, the permutation is the last permutation.
+                return
+
+            # 2. Find the largest index j greater than i such that a[i] < a[j].
+            for j in range(size - 1, i, -1):
+                if current[i] < current[j]:
+                    break
+
+            else:  # pragma: NO COVER
+                raise RuntimeError("Broken algorithm")
+
+            # 3. Swap the value of a[i] with that of a[j].
+            temp = current[i]
+            current[i] = current[j]
+            current[j] = temp
+
+            # 4. Reverse the sequence from a[i + 1] up to and including the final
+            # element a[n].
+            current = current[: i + 1] + list(reversed(current[i + 1 :]))

test_utils/prefixer.py

@@ -0,0 +1,82 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import random
+import re
+
+from typing import Union
+
+_RESOURCE_DATE_FORMAT = "%Y%m%d%H%M%S"
+_RESOURCE_DATE_LENGTH = 4 + 2 + 2 + 2 + 2 + 2
+_RE_SEPARATORS = re.compile(r"[/\-\\_]")
+
+
+def _common_prefix(repo, relative_dir, separator="_"):
+    repo = _RE_SEPARATORS.sub(separator, repo)
+    relative_dir = _RE_SEPARATORS.sub(separator, relative_dir)
+    return f"{repo}{separator}{relative_dir}"
+
+
+class Prefixer(object):
+    """Create/manage resource IDs for system testing.
+
+    Usage:
+
+    Creating resources:
+
+    >>> import test_utils.prefixer
+    >>> prefixer = test_utils.prefixer.Prefixer("python-bigquery", "samples/snippets")
+    >>> dataset_id = prefixer.create_prefix() + "my_sample"
+
+    Cleaning up resources:
+
+    >>> @pytest.fixture(scope="session", autouse=True)
+    ... def cleanup_datasets(bigquery_client: bigquery.Client):
+    ...     for dataset in bigquery_client.list_datasets():
+    ...         if prefixer.should_cleanup(dataset.dataset_id):
+    ...             bigquery_client.delete_dataset(
+    ...                 dataset, delete_contents=True, not_found_ok=True
+    """
+
+    def __init__(
+        self, repo, relative_dir, separator="_", cleanup_age=datetime.timedelta(days=1)
+    ):
+        self._separator = separator
+        self._cleanup_age = cleanup_age
+        self._prefix = _common_prefix(repo, relative_dir, separator=separator)
+
+    def create_prefix(self) -> str:
+        now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+        timestamp = now.strftime(_RESOURCE_DATE_FORMAT)
+        random_string = hex(random.randrange(0x1000000))[2:]
+        return f"{self._prefix}{self._separator}{timestamp}{self._separator}{random_string}"
+
+    def _name_to_date(self, resource_name: str) -> Union[datetime.datetime, None]:
+        start_date = len(self._prefix) + len(self._separator)
+        date_string = resource_name[start_date : start_date + _RESOURCE_DATE_LENGTH]
+        try:
+            parsed_date = datetime.datetime.strptime(date_string, _RESOURCE_DATE_FORMAT)
+            return parsed_date
+        except ValueError:
+            return None
+
+    def should_cleanup(self, resource_name: str) -> bool:
+        now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+        yesterday = now - self._cleanup_age
+        if not resource_name.startswith(self._prefix):
+            return False
+
+        created_date = self._name_to_date(resource_name)
+        return created_date is not None and created_date < yesterday

test_utils/py.typed

@@ -0,0 +1,2 @@
+# Marker file for PEP 561.
+# The test_utils package uses inline types.