<commit_before>#!/usr/bin/env vpython
# Copyright 2016 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

from __future__ import print_function
from future.utils import iteritems

import contextlib
import json
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time

from parameterized import parameterized, parameterized_class

import test_env

from recipe_engine.internal.engine import _shell_quote

from recipe_engine.third_party import luci_context

# prevent LUCI_CONTEXT leakage :)
os.environ.pop(luci_context.ENV_KEY, None)

class RunTest(test_env.RecipeEngineUnitTest):
  @parameterized.expand(
    [('py2', False), ('py3', True)],
    name_func=lambda func, _, param: '%s_%s' % (func.__name__, param.args[0]),
  )
  def test_run(self, _, py3):
    deps = self.FakeRecipeDeps()

    with deps.main_repo.write_module('mod') as mod:
      mod.api.write('''
      def do_thing(self):
        self.m.step('do the thing', ['echo', 'thing'])
      ''')
    with deps.main_repo.write_recipe('my_recipe') as recipe:
      recipe.DEPS = ['mod']
      recipe.RunSteps.write('''
        api.mod.do_thing()
      ''')
      recipe.GenTests.write('pass')

    output, retcode = deps.main_repo.recipes_py('-v', '-v', 'run', 'my_recipe',
                                                py3=py3)
    self.assertEqual(retcode, 0,
                     'ret code is not zero. Recipe output\n%s' % output)

  def test_run_incomplete_deps(self):
    deps = self.FakeRecipeDeps()

    a = deps.add_repo('a', detached=True)
    a.commit('a something')

    b = deps.add_repo('b', detached=True)
    b.commit('b something')

    a.add_dep('b')
    a.commit('add b dep')

    deps.main_repo.add_dep('a')
    a.commit('add a dep')

    output, retcode = deps.main_repo.recipes_py('fetch')
    self.assertEqual(retcode, 1)
    self.assertIn('Repo \'a\' depends on [\'b\'], which is missing', output)

  def test_run_circular_deps(self):
    deps = self.FakeRecipeDeps()

    a = deps.add_repo('a', detached=True)
    a.commit('a something')

    a.add_dep('main')
    a.commit('add main dep')

    deps.main_repo.add_dep('a')
    a.commit('add a dep')

    output, retcode = deps.main_repo.recipes_py('fetch')
    self.assertEqual(retcode, 1)
    self.assertIn(
        'Dependency \'a\' has circular dependency on \'main\'', output)

class RunSmokeTest(test_env.RecipeEngineUnitTest):
  def _run_cmd(self, recipe, properties=None, engine_args=()):
    script_path = os.path.join(test_env.ROOT_DIR, 'recipes.py')

    proplist = [
      '%s=%s' % (k, json.dumps(v)) for k, v in iteritems(properties or {})
    ]

    return (
      ['python', script_path] +
      list(engine_args) +
      ['run', recipe] +
      proplist
    )

  def _test_recipe(self, recipe, properties=None, env=None):
    workdir = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    try:
      proc = subprocess.Popen(
          self._run_cmd(recipe, properties),
          cwd=workdir,
          stdout=subprocess.PIPE,
          stderr=subprocess.STDOUT,
          env=env)
      stdout = proc.communicate()
      self.assertEqual(0, proc.returncode, '%d != %d when testing %s:\n%s' % (
          0, proc.returncode, recipe, stdout))
    finally:
      shutil.rmtree(workdir, ignore_errors=True)

  def test_examples(self):
    env = os.environ.copy()

    # Set the "RECIPE_ENGINE_CONTEXT_TEST" environment variable to a known
    # value, "default". This is used by the "context:tests/env" recipe module
    # as a basis for runtime tests.
    env['RECIPE_ENGINE_CONTEXT_TEST'] = 'default'

    tests = [
      ['context:examples/full'],
      ['context:tests/env'],
      ['file:examples/copy'],
      ['file:examples/copytree'],
      ['file:examples/glob'],
      ['futures:examples/lazy_fan_out_in'],
      ['json:examples/full'],
      ['path:examples/full'],
      ['python:examples/full'],
      ['raw_io:examples/full'],
      ['step:examples/full'],

      ['engine_tests/functools_partial'],
    ]
    for test in tests:
      print("running", test)
      self._test_recipe(*test, env=env)

  def test_bad_subprocess(self):
    now = time.time()
    self._test_recipe('engine_tests/bad_subprocess')
    after = time.time()

    # Test has a daemon that holds on to stdout for 30s, but the daemon's parent
    # process (e.g. the one that recipe engine actually runs) quits immediately.
    # If this takes longer than 10 seconds to run (there can be overhead in
    # running the engine/cipd/protoc/etc.), we consider it failed.
    self.assertLess(after - now, 10)

  def test_shell_quote(self):
    # For regular-looking commands we shouldn't need any specialness.
    self.assertEqual(
        _shell_quote('/usr/bin/python-wrapper.bin'),
        '/usr/bin/python-wrapper.bin')

    STRINGS = [
        'Simple.Command123/run',
        'Command with spaces',
        'Command with "quotes"',
        "I have 'single quotes'",
        'Some \\Esc\ape Seque\nces/',
        u'Unicode makes me \u2609\u203f\u2299'.encode('utf-8'),
    ]

    for s in STRINGS:
      quoted = _shell_quote(s)

      # We shouldn't ever get an actual newline in a command, that's awful
      # for copypasta.
      self.assertNotRegexpMatches(quoted, '\n')

      # We should be able to paste any argument into bash & zsh and get
      # exactly what subprocess did.
      bash_output = subprocess.check_output([
          'bash', '-c', '/bin/echo %s' % quoted])
      self.assertEqual(bash_output, s + '\n')

      # zsh is untested because zsh isn't provisioned on our bots.
      # zsh_output = subprocess.check_output([
      #     'zsh', '-c', '/bin/echo %s' % quoted])
      # self.assertEqual(zsh_output.decode('utf-8'), s + '\n')

@parameterized_class(
  [{"py_version": 2}, {"py_version": 3}],
  class_name_func=(
    lambda cls, _, params: '%s_PY%d' % (cls.__name__, params['py_version'])),
)
class LuciexeSmokeTest(test_env.RecipeEngineUnitTest):
  def _wait_for_file(self, filename, duration):
    begin = time.time()
    while True:
      self.assertLessEqual(time.time() - begin, duration,
                           'took too long to find ' + filename)
      try:
        with open(filename, 'r') as f:
          return f.read()
      except IOError:
        time.sleep(.5)

  @contextlib.contextmanager
  def _run_bbagent(self, properties, grace_period=30):
    workdir = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    proc = None
    try:
      pidfile = os.path.join(workdir, 'pidfile')
      fake_bbagent = os.path.join(test_env.ROOT_DIR, 'misc', 'fake_bbagent.sh')

      env = os.environ.copy()
      env.pop(luci_context.ENV_KEY, None)
      env['WD'] = workdir
      env['LUCI_GRACE_PERIOD'] = str(grace_period)
      if self.py_version == 3:
        env['RECIPES_USE_PY3'] = 'true'

      proc = subprocess.Popen(
          [fake_bbagent, "--pid-file", pidfile],
          stdin=subprocess.PIPE,
          env=env)
      proc.stdin.write(json.dumps({
        "input": {
          "properties": properties,
        },
      }))
      proc.stdin.close()

      engine_pid = int(self._wait_for_file(pidfile, 30).strip())

      yield proc, engine_pid
    finally:
      if proc and proc.poll() is None:
        proc.kill()
      shutil.rmtree(workdir, ignore_errors=True)

  def _test_bbagent(self, properties, grace_period=30, timeout=30):
    scrap = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    try:
      wd = os.path.join(scrap, 'wd')
      os.mkdir(wd)

      outfile = os.path.join(scrap, 'final_build.json')
      pidfile = os.path.join(wd, 'pidfile')
      fake_bbagent = os.path.join(test_env.ROOT_DIR, 'misc', 'fake_bbagent.sh')

      deadline = time.time() + timeout

      env = os.environ.copy()
      env.pop(luci_context.ENV_KEY, None)
      env['WD'] = wd
      env['LUCI_GRACE_PERIOD'] = str(grace_period)
      env['LUCI_SOFT_DEADLINE'] = str(deadline)
      env['FAKE_BBAGENT_OUTFILE'] = outfile
      if self.py_version == 3:
        env['RECIPES_USE_PY3'] = 'true'
      proc = subprocess.Popen(
          [fake_bbagent, "--pid-file", pidfile],
          stdin=subprocess.PIPE,
          env=env)

      proc.stdin.write(json.dumps({
        "input": {
          "properties": properties,
        },
      }))
      proc.stdin.close()

      engine_pid = int(self._wait_for_file(pidfile, 30).strip())
      did_soft_deadline = False

      while True:
        if proc.poll() is not None:
          with open(os.path.join(wd, 'logs', 'stderr')) as log:
            print()
            print("Raw engine logs:")
            sys.stdout.write(log.read())
          with open(outfile) as of:
            bpdata = of.read()
            print()
            print("Final build.proto:")
            sys.stdout.write(bpdata)
            return json.loads(bpdata)

        if deadline and time.time() > deadline:
          if not did_soft_deadline:
            print("Hit soft deadline")
            did_soft_deadline = True
            os.kill(engine_pid, signal.SIGTERM)
            deadline += grace_period
          else:
            print("Hit hard deadline")
            os.kill(engine_pid, signal.SIGKILL)
            deadline = None

        time.sleep(1)

    finally:
      shutil.rmtree(scrap, ignore_errors=True)

  def test_early_terminate(self):
    scrap = tempfile.mkdtemp(prefix='recipe_engine-run_test-scrap-')
    try:
      # The recipe will make a bunch of subprocesses which whill touch this
      # every ~second.
      output_touchfile = os.path.join(scrap, 'output_touchfile')
      running_touchfile = os.path.join(scrap, 'running_touchfile')

      props = {
        'recipe': 'engine_tests/early_termination',
        'output_touchfile': output_touchfile,
        'running_touchfile': running_touchfile,
      }
      with self._run_bbagent(props, grace_period=5) as (proc, engine_pid):
        # Wait up to 20s for the recipe to indicate that it launched all its
        # subprocesses.
        self._wait_for_file(running_touchfile, 20)

        # Ok, the recipe is all up and running now. Let's give the command
        # a poke and wait for a bit.
        os.kill(engine_pid, signal.SIGTERM)

        # from this point the recipe should teardown in ~5s. We give it 10 to
        # be generous.
        time.sleep(10)
        self.assertIsNotNone(proc.poll())

        # sample the output_touchfile
        mtime = os.stat(output_touchfile).st_mtime
        # now wait a bit to see if anything's still touching it
        time.sleep(5)
        self.assertEqual(mtime, os.stat(output_touchfile).st_mtime)
    finally:
      shutil.rmtree(scrap, ignore_errors=True)

  def test_add_tags(self):
    final_build = self._test_bbagent(
        {'recipe': 'buildbucket:tests/add_tags'},
    )
    self.assertListEqual(
      sorted(final_build['tags'], key=lambda tag: (tag['key'], tag['value'])),
      [
        {'key': 'hide-in-gerrit', 'value': 'pointless'},
        {'key': 'k1', 'value': 'v1'},
        {'key': 'k2', 'value': 'v2'},
        {'key': 'k2', 'value': 'v2_1'},
      ],
    )

  def test_output_gitiles(self):
    final_build = self._test_bbagent(
        {'recipe': 'buildbucket:tests/output_commit'},
    )
    self.assertDictEqual(final_build['output']['gitiles_commit'], {
      'host': 'chromium.googlesource.com',
      'id': 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa',
      'position': 42,
      'project': 'infra/infra',
      'ref': 'refs/heads/main',
    })

  def test_proto_output_properties(self):
    final_build = self._test_bbagent(
        {'recipe': 'engine_tests/proto_output_properties'},
    )
    output_props = final_build['output']['properties']
    self.assertDictEqual(output_props['$mod/proto_out'], {
      'str': 'foo',
      'strs': ['bar', 'baz'],
      'msg' : {
        'num': 1,
        'nums': [10, 11, 12],
      },
    })

  def test_external_timeout(self):
    final_build = self._test_bbagent(
        {'recipe': 'engine_tests/long_sleep'},
        timeout=10,
    )
    for step in final_build['steps']:
      del step['end_time']
      del step['logs']
      del step['start_time']
      if step['name'] == 'setup_build':
        del step['summary_markdown']
    self.assertDictEqual(final_build, {
      'status': 'INFRA_FAILURE',
      'summary_markdown': (
        "Infra Failure: Step('sleep forever') (canceled) (retcode: -15)"),
      'steps': [
        {'name': 'setup_build', 'status': 'SUCCESS'},
        {'name': 'sleep a bit', 'status': 'FAILURE'},
        {'name': 'sleep forever', 'status': 'CANCELED'},
      ],
    })

  def test_external_timeout_recovery(self):
    final_build = self._test_bbagent(
        {
          'recipe': 'engine_tests/long_sleep',
          'recover': True,
        },
        timeout=10,
    )
    self.assertEqual(final_build['status'], 'SUCCESS')

  def test_nonexistent_command(self):
    final_build = self._test_bbagent(
      {'recipe': 'engine_tests/nonexistent_command'},
    )
    self.assertEqual(final_build['status'], 'SUCCESS')

if __name__ == '__main__':
  test_env.main()
<commit_msg>Recipe 'run' command has a weirdo concept of how working directory is handled.<commit_after>#!/usr/bin/env vpython
# Copyright 2016 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

from __future__ import print_function
from future.utils import iteritems

import contextlib
import json
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time

from parameterized import parameterized, parameterized_class

import test_env

from recipe_engine.internal.engine import _shell_quote

from recipe_engine.third_party import luci_context

# prevent LUCI_CONTEXT leakage :)
os.environ.pop(luci_context.ENV_KEY, None)

class RunTest(test_env.RecipeEngineUnitTest):
  @parameterized.expand(
    [('py2', False), ('py3', True)],
    name_func=lambda func, _, param: '%s_%s' % (func.__name__, param.args[0]),
  )
  def test_run(self, _, py3):
    deps = self.FakeRecipeDeps()

    with deps.main_repo.write_module('mod') as mod:
      mod.api.write('''
      def do_thing(self):
        self.m.step('do the thing', ['echo', 'thing'])
      ''')
    with deps.main_repo.write_recipe('my_recipe') as recipe:
      recipe.DEPS = ['mod']
      recipe.RunSteps.write('''
        api.mod.do_thing()
      ''')
      recipe.GenTests.write('pass')

    output, retcode = deps.main_repo.recipes_py('-v', '-v', 'run', 'my_recipe',
                                                py3=py3)
    self.assertEqual(retcode, 0,
                     'ret code is not zero. Recipe output\n%s' % output)

  def test_run_incomplete_deps(self):
    deps = self.FakeRecipeDeps()

    a = deps.add_repo('a', detached=True)
    a.commit('a something')

    b = deps.add_repo('b', detached=True)
    b.commit('b something')

    a.add_dep('b')
    a.commit('add b dep')

    deps.main_repo.add_dep('a')
    a.commit('add a dep')

    output, retcode = deps.main_repo.recipes_py('fetch')
    self.assertEqual(retcode, 1)
    self.assertIn('Repo \'a\' depends on [\'b\'], which is missing', output)

  def test_run_circular_deps(self):
    deps = self.FakeRecipeDeps()

    a = deps.add_repo('a', detached=True)
    a.commit('a something')

    a.add_dep('main')
    a.commit('add main dep')

    deps.main_repo.add_dep('a')
    a.commit('add a dep')

    output, retcode = deps.main_repo.recipes_py('fetch')
    self.assertEqual(retcode, 1)
    self.assertIn(
        'Dependency \'a\' has circular dependency on \'main\'', output)

class RunSmokeTest(test_env.RecipeEngineUnitTest):
  def _run_cmd(self, recipe, workdir, properties=None, engine_args=()):
    script_path = os.path.join(test_env.ROOT_DIR, 'recipes.py')

    proplist = [
      '%s=%s' % (k, json.dumps(v)) for k, v in iteritems(properties or {})
    ]

    return (
      ['python', script_path] +
      list(engine_args) +
      ['run', '--workdir', workdir, recipe] +
      proplist
    )

  def _test_recipe(self, recipe, properties=None, env=None):
    workdir = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    try:
      proc = subprocess.Popen(
          self._run_cmd(recipe, workdir, properties),
          stdout=subprocess.PIPE,
          stderr=subprocess.STDOUT,
          env=env)
      stdout = proc.communicate()
      self.assertEqual(0, proc.returncode, '%d != %d when testing %s:\n%s' % (
          0, proc.returncode, recipe, stdout))
    finally:
      shutil.rmtree(workdir, ignore_errors=True)

  def test_examples(self):
    env = os.environ.copy()

    # Set the "RECIPE_ENGINE_CONTEXT_TEST" environment variable to a known
    # value, "default". This is used by the "context:tests/env" recipe module
    # as a basis for runtime tests.
    env['RECIPE_ENGINE_CONTEXT_TEST'] = 'default'

    tests = [
      ['context:examples/full'],
      ['context:tests/env'],
      ['file:examples/copy'],
      ['file:examples/copytree'],
      ['file:examples/glob'],
      ['futures:examples/lazy_fan_out_in'],
      ['json:examples/full'],
      ['path:examples/full'],
      ['python:examples/full'],
      ['raw_io:examples/full'],
      ['step:examples/full'],

      ['engine_tests/functools_partial'],
    ]
    for test in tests:
      print("running", test)
      self._test_recipe(*test, env=env)

  def test_bad_subprocess(self):
    now = time.time()
    self._test_recipe('engine_tests/bad_subprocess')
    after = time.time()

    # Test has a daemon that holds on to stdout for 30s, but the daemon's parent
    # process (e.g. the one that recipe engine actually runs) quits immediately.
    # If this takes longer than 10 seconds to run (there can be overhead in
    # running the engine/cipd/protoc/etc.), we consider it failed.
    self.assertLess(after - now, 10)

  def test_shell_quote(self):
    # For regular-looking commands we shouldn't need any specialness.
    self.assertEqual(
        _shell_quote('/usr/bin/python-wrapper.bin'),
        '/usr/bin/python-wrapper.bin')

    STRINGS = [
        'Simple.Command123/run',
        'Command with spaces',
        'Command with "quotes"',
        "I have 'single quotes'",
        'Some \\Esc\ape Seque\nces/',
        u'Unicode makes me \u2609\u203f\u2299'.encode('utf-8'),
    ]

    for s in STRINGS:
      quoted = _shell_quote(s)

      # We shouldn't ever get an actual newline in a command, that's awful
      # for copypasta.
      self.assertNotRegexpMatches(quoted, '\n')

      # We should be able to paste any argument into bash & zsh and get
      # exactly what subprocess did.
      bash_output = subprocess.check_output([
          'bash', '-c', '/bin/echo %s' % quoted])
      self.assertEqual(bash_output, s + '\n')

      # zsh is untested because zsh isn't provisioned on our bots.
      # zsh_output = subprocess.check_output([
      #     'zsh', '-c', '/bin/echo %s' % quoted])
      # self.assertEqual(zsh_output.decode('utf-8'), s + '\n')

@parameterized_class(
  [{"py_version": 2}, {"py_version": 3}],
  class_name_func=(
    lambda cls, _, params: '%s_PY%d' % (cls.__name__, params['py_version'])),
)
class LuciexeSmokeTest(test_env.RecipeEngineUnitTest):
  def _wait_for_file(self, filename, duration):
    begin = time.time()
    while True:
      self.assertLessEqual(time.time() - begin, duration,
                           'took too long to find ' + filename)
      try:
        with open(filename, 'r') as f:
          return f.read()
      except IOError:
        time.sleep(.5)

  @contextlib.contextmanager
  def _run_bbagent(self, properties, grace_period=30):
    workdir = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    proc = None
    try:
      pidfile = os.path.join(workdir, 'pidfile')
      fake_bbagent = os.path.join(test_env.ROOT_DIR, 'misc', 'fake_bbagent.sh')

      env = os.environ.copy()
      env.pop(luci_context.ENV_KEY, None)
      env['WD'] = workdir
      env['LUCI_GRACE_PERIOD'] = str(grace_period)
      if self.py_version == 3:
        env['RECIPES_USE_PY3'] = 'true'

      proc = subprocess.Popen(
          [fake_bbagent, "--pid-file", pidfile],
          stdin=subprocess.PIPE,
          env=env)
      proc.stdin.write(json.dumps({
        "input": {
          "properties": properties,
        },
      }))
      proc.stdin.close()

      engine_pid = int(self._wait_for_file(pidfile, 30).strip())

      yield proc, engine_pid
    finally:
      if proc and proc.poll() is None:
        proc.kill()
      shutil.rmtree(workdir, ignore_errors=True)

  def _test_bbagent(self, properties, grace_period=30, timeout=30):
    scrap = tempfile.mkdtemp(prefix='recipe_engine_run_test-')
    try:
      wd = os.path.join(scrap, 'wd')
      os.mkdir(wd)

      outfile = os.path.join(scrap, 'final_build.json')
      pidfile = os.path.join(wd, 'pidfile')
      fake_bbagent = os.path.join(test_env.ROOT_DIR, 'misc', 'fake_bbagent.sh')

      deadline = time.time() + timeout

      env = os.environ.copy()
      env.pop(luci_context.ENV_KEY, None)
      env['WD'] = wd
      env['LUCI_GRACE_PERIOD'] = str(grace_period)
      env['LUCI_SOFT_DEADLINE'] = str(deadline)
      env['FAKE_BBAGENT_OUTFILE'] = outfile
      if self.py_version == 3:
        env['RECIPES_USE_PY3'] = 'true'
      proc = subprocess.Popen(
          [fake_bbagent, "--pid-file", pidfile],
          stdin=subprocess.PIPE,
          env=env)

      proc.stdin.write(json.dumps({
        "input": {
          "properties": properties,
        },
      }))
      proc.stdin.close()

      engine_pid = int(self._wait_for_file(pidfile, 30).strip())
      did_soft_deadline = False

      while True:
        if proc.poll() is not None:
          with open(os.path.join(wd, 'logs', 'stderr')) as log:
            print()
            print("Raw engine logs:")
            sys.stdout.write(log.read())
          with open(outfile) as of:
            bpdata = of.read()
            print()
            print("Final build.proto:")
            sys.stdout.write(bpdata)
            return json.loads(bpdata)

        if deadline and time.time() > deadline:
          if not did_soft_deadline:
            print("Hit soft deadline")
            did_soft_deadline = True
            os.kill(engine_pid, signal.SIGTERM)
            deadline += grace_period
          else:
            print("Hit hard deadline")
            os.kill(engine_pid, signal.SIGKILL)
            deadline = None

        time.sleep(1)

    finally:
      shutil.rmtree(scrap, ignore_errors=True)

  def test_early_terminate(self):
    scrap = tempfile.mkdtemp(prefix='recipe_engine-run_test-scrap-')
    try:
      # The recipe will make a bunch of subprocesses which whill touch this
      # every ~second.
      output_touchfile = os.path.join(scrap, 'output_touchfile')
      running_touchfile = os.path.join(scrap, 'running_touchfile')

      props = {
        'recipe': 'engine_tests/early_termination',
        'output_touchfile': output_touchfile,
        'running_touchfile': running_touchfile,
      }
      with self._run_bbagent(props, grace_period=5) as (proc, engine_pid):
        # Wait up to 20s for the recipe to indicate that it launched all its
        # subprocesses.
        self._wait_for_file(running_touchfile, 20)

        # Ok, the recipe is all up and running now. Let's give the command
        # a poke and wait for a bit.
        os.kill(engine_pid, signal.SIGTERM)

        # from this point the recipe should teardown in ~5s. We give it 10 to
        # be generous.
        time.sleep(10)
        self.assertIsNotNone(proc.poll())

        # sample the output_touchfile
        mtime = os.stat(output_touchfile).st_mtime
        # now wait a bit to see if anything's still touching it
        time.sleep(5)
        self.assertEqual(mtime, os.stat(output_touchfile).st_mtime)
    finally:
      shutil.rmtree(scrap, ignore_errors=True)

  def test_add_tags(self):
    final_build = self._test_bbagent(
        {'recipe': 'buildbucket:tests/add_tags'},
    )
    self.assertListEqual(
      sorted(final_build['tags'], key=lambda tag: (tag['key'], tag['value'])),
      [
        {'key': 'hide-in-gerrit', 'value': 'pointless'},
        {'key': 'k1', 'value': 'v1'},
        {'key': 'k2', 'value': 'v2'},
        {'key': 'k2', 'value': 'v2_1'},
      ],
    )

  def test_output_gitiles(self):
    final_build = self._test_bbagent(
        {'recipe': 'buildbucket:tests/output_commit'},
    )
    self.assertDictEqual(final_build['output']['gitiles_commit'], {
      'host': 'chromium.googlesource.com',
      'id': 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa',
      'position': 42,
      'project': 'infra/infra',
      'ref': 'refs/heads/main',
    })

  def test_proto_output_properties(self):
    final_build = self._test_bbagent(
        {'recipe': 'engine_tests/proto_output_properties'},
    )
    output_props = final_build['output']['properties']
    self.assertDictEqual(output_props['$mod/proto_out'], {
      'str': 'foo',
      'strs': ['bar', 'baz'],
      'msg' : {
        'num': 1,
        'nums': [10, 11, 12],
      },
    })

  def test_external_timeout(self):
    final_build = self._test_bbagent(
        {'recipe': 'engine_tests/long_sleep'},
        timeout=10,
    )
    for step in final_build['steps']:
      del step['end_time']
      del step['logs']
      del step['start_time']
      if step['name'] == 'setup_build':
        del step['summary_markdown']
    self.assertDictEqual(final_build, {
      'status': 'INFRA_FAILURE',
      'summary_markdown': (
        "Infra Failure: Step('sleep forever') (canceled) (retcode: -15)"),
      'steps': [
        {'name': 'setup_build', 'status': 'SUCCESS'},
        {'name': 'sleep a bit', 'status': 'FAILURE'},
        {'name': 'sleep forever', 'status': 'CANCELED'},
      ],
    })

  def test_external_timeout_recovery(self):
    final_build = self._test_bbagent(
        {
          'recipe': 'engine_tests/long_sleep',
          'recover': True,
        },
        timeout=10,
    )
    self.assertEqual(final_build['status'], 'SUCCESS')

  def test_nonexistent_command(self):
    final_build = self._test_bbagent(
      {'recipe': 'engine_tests/nonexistent_command'},
    )
    self.assertEqual(final_build['status'], 'SUCCESS')

if __name__ == '__main__':
  test_env.main()
<|endoftext|>


<commit_before>from __future__ import absolute_import
import copy, threading
from collections import namedtuple
from .interfaces import Cursor, Static, Journal, Memory, Change, CannotCommit
from .log import log, weaklog

__all__ = (
    'memory', 'journal',
    'alloc', 'dealloc', 'read_unsaved', 'read_saved', 'commit_changes', 'write',
    'save_state', 'revert_state', 'change_state'
)

### Generic Operations

def alloc(journal, cursor, state):
    journal.allocate(cursor, state)
    return cursor

def dealloc(journal, cursor):
    journal.delete(cursor)

def read_unsaved(journal, cursor):
    return good_value(cursor, journal.read_unsaved(cursor))

def read_saved(journal, cursor, *default):
    return good_value(cursor, journal.read_saved(cursor), *default)

def commit_changes(source, nested):
    if source is nested:
        raise RuntimeError("A journal can't be committed to itself.")
    source.commit_changes(nested)

def write(journal, cursor):
    return good_value(cursor, journal.write(cursor))

def save_state(journal, cursor, force=False):
    return journal.save_state(cursor, force)

def revert_state(journal, cursor):
    journal.revert_state(cursor)

def change_state(op, journal, what, *args, **kwargs):
    if isinstance(what, Cursor):
        op(journal, what, *args, **kwargs)
    else:
        ## cache cursors in a list so the log can be modified.
        for cursor in list(what):
            op(journal, cursor, *args, **kwargs)

### Journals

class change(namedtuple('changes', 'cursor orig state'), Change):
    pass

class journal(Journal):
    LogType = log
    name = None
    source = None

    def __init__(self, name, source):
        self.name = name
        self.source = source
        self.read_log = self.LogType()
        self.commit_log = self.LogType()
        self.write_log = self.LogType()

        ## Aggressively notify even though there's no activity yet.
        ## This simplifies weird situations like insert-only or
        ## delete-only transactions.
        self.begun = False; self.notify()

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, str(self))

    def __str__(self):
        return self.name

    def begin(self, nested):
        self.notify()

    def notify(self):
        if not self.begun:
            self.source.begin(self)
            self.begun = True

    def committed(self):
        self.begun = False

    def make_journal(self, name):
        return type(self)(name, self)

    def allocate(self, cursor, state):
        self.commit_log.allocate(cursor, state)

    def read_unsaved(self, cursor):
        try:
            return self.write_log[cursor]
        except KeyError:
            return self.read_saved(cursor)

    def read_saved(self, cursor):
        try:
            return self.commit_log[cursor]
        except KeyError:
            return self._read(cursor)

    def _read(self, cursor):
        try:
            return self.read_log[cursor]
        except KeyError:
            self.notify()
            return log_read(self, cursor)

    def write(self, cursor):
        try:
            return self.write_log[cursor]
        except:
            return log_write(self, cursor)

    def delete(self, cursor):
        self._read(cursor) ## Make sure this has been read in
        log_delete(self, cursor)

    def save_state(self, cursor, force=False):
        try:
            state = self.write_log.pop(cursor)
        except KeyError:
            ## There's nothing to write; just return.
            if force:
                self.commit_log[cursor] = self.read_saved(cursor)
                return True
            else:
                return False
        else:
            self.update_commit_log(cursor, state)
            return True

    def revert_state(self, cursor):
        self.write_log.pop(cursor, None)

    def commit_changes(self, nested):
        ## A journal is single-threaded; state can be blindly copied
        ## in.
        for (cursor, orig, state) in nested.changed():
            self.update_commit_log(cursor, state)
        nested.committed()

    def unsaved(self):
        return (
            k for (k, v) in self.write_log
            if read_saved(self, k, None) is not v
        )

    def changed(self):
        return (
            change(k, get_state(self.read_log, k), v)
            for (k, v) in self.commit_log
        )

    def read(self):
        return iter(self.read_log)

    def update_read_log(self, cursor, state):
        self.read_log[cursor] = state

    def update_write_log(self, cursor, state):
        self.write_log[cursor] = state

    def update_commit_log(self, cursor, state):
        self.commit_log[cursor] = copy_state(state)

class memory(Memory):
    JournalType = journal
    LogType = weaklog
    name = None

    def __init__(self, name='*memory*', check_read=True, check_write=True):
        self.name = name
        self.write_lock = threading.RLock()
        self.check_read = check_read
        self.check_write = check_write
        self.mem = self.LogType()

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, str(self))

    def __str__(self):
        return self.name

    def __contains__(self, cursor):
        return cursor in self.mem

    def begin(self, nested):
        pass

    def committed(self):
        pass

    def make_journal(self, name):
        return self.JournalType(name, self)

    def allocate(self, cursor, state):
        self.mem.allocate(cursor, state)

    def read_saved(self, cursor):
        return self.mem[cursor]

    def commit_changes(self, nested):
        self.with_changes(nested, nested.changed(), self.write_changes)

    def with_changes(self, trans, changed, op):
        with self.write_lock:
            self.verify_read(trans.read())
            op(trans, self.verify_write(changed))
            trans.committed()

    def verify_read(self, read):
        if self.check_read:
            verify_read(self.mem, read)

    def verify_write(self, changed):
        if self.check_write:
            return verify_write(self.mem, changed)
        else:
            return unverified_write(changed)

    def write_changes(self, nested, changed):
        for (cursor, state) in changed:
            if is_deleted(state):
                self.expire(cursor)
            else:
                self.mem[cursor] = state

    def expire(self, cursor):
        if isinstance(cursor, Cursor):
            self.mem.pop(cursor, None)
        elif callable(cursor):
            for cur in self.mem.keys():
                if cursor(cur):
                    del self.mem[cur]
        else:
            raise TypeError('expecting cursor or callable', cursor)

    def expire_all(self, cursors=None):
        if cursors:
            for cursor in cursors:
                self.expire(cursor)
        else:
            self.mem.clear()

### State

copy_state = copy.deepcopy

class sentinal(object):
    __slots__ = ('name')

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, self.name)

    def __copy__(self):
        return self

    def __deepcopy__(self, memo):
        return self

INSERTED = sentinal('inserted')
DELETED = sentinal('deleted')

def good_value(cursor, value, *default):
    if value is INSERTED or value is DELETED:
        if default:
            return default[0]
        else:
            raise ValueError, id(cursor)
    else:
        return value

def is_deleted(obj):
    return obj is DELETED

def is_inserted(obj):
    return obj is INSERTED

### Logged Operations

def log_read(journal, cursor):
    state = read_saved(journal.source, cursor, INSERTED)
    journal.update_read_log(cursor, state)
    return state

def log_write(journal, cursor):
    if isinstance(cursor.__id__, Static):
        raise TypeError('static cursors are not writable', cursor)
    state = copy_state(read_saved(journal, cursor))
    journal.update_write_log(cursor, state)
    return state

def log_delete(journal, cursor):
    if isinstance(cursor.__id__, Static):
        raise TypeError('static cursors are not writable', cursor)
    journal.update_write_log(cursor, DELETED)

### Operations on Logs

def get_state(log, cursor):
    return log.get(cursor, INSERTED)

def verify_read(log, read):
    conflicts = [(c, s) for (c, s) in read if log.get(c) != s]
    if conflicts:
        raise CannotCommit, conflicts

def verify_write(log, changed):
    changed, conflicts = partition_conflicts(log, changed)
    if conflicts:
        raise CannotCommit, conflicts
    return changed

def unverified_write(changed):
    return list((c, s) for (c, o, s) in changed)

def partition_conflicts(log, changed):
    good = []; bad = []
    for (cursor, orig, state) in changed:
        current = get_state(log, cursor)
        (good if current is orig else bad).append((cursor, state))
    return good, bad
<commit_msg>rename _read<commit_after>from __future__ import absolute_import
import copy, threading
from collections import namedtuple
from .interfaces import Cursor, Static, Journal, Memory, Change, CannotCommit
from .log import log, weaklog

__all__ = (
    'memory', 'journal',
    'alloc', 'dealloc', 'read_unsaved', 'read_saved', 'commit_changes', 'write',
    'save_state', 'revert_state', 'change_state'
)

### Generic Operations

def alloc(journal, cursor, state):
    journal.allocate(cursor, state)
    return cursor

def dealloc(journal, cursor):
    journal.delete(cursor)

def read_unsaved(journal, cursor):
    return good_value(cursor, journal.read_unsaved(cursor))

def read_saved(journal, cursor, *default):
    return good_value(cursor, journal.read_saved(cursor), *default)

def commit_changes(source, nested):
    if source is nested:
        raise RuntimeError("A journal can't be committed to itself.")
    source.commit_changes(nested)

def write(journal, cursor):
    return good_value(cursor, journal.write(cursor))

def save_state(journal, cursor, force=False):
    return journal.save_state(cursor, force)

def revert_state(journal, cursor):
    journal.revert_state(cursor)

def change_state(op, journal, what, *args, **kwargs):
    if isinstance(what, Cursor):
        op(journal, what, *args, **kwargs)
    else:
        ## cache cursors in a list so the log can be modified.
        for cursor in list(what):
            op(journal, cursor, *args, **kwargs)

### Journals

class change(namedtuple('changes', 'cursor orig state'), Change):
    pass

class journal(Journal):
    LogType = log
    name = None
    source = None

    def __init__(self, name, source):
        self.name = name
        self.source = source
        self.read_log = self.LogType()
        self.commit_log = self.LogType()
        self.write_log = self.LogType()

        ## Aggressively notify even though there's no activity yet.
        ## This simplifies weird situations like insert-only or
        ## delete-only transactions.
        self.begun = False; self.notify()

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, str(self))

    def __str__(self):
        return self.name

    def begin(self, nested):
        self.notify()

    def notify(self):
        if not self.begun:
            self.source.begin(self)
            self.begun = True

    def committed(self):
        self.begun = False

    def make_journal(self, name):
        return type(self)(name, self)

    def allocate(self, cursor, state):
        self.commit_log.allocate(cursor, state)

    def read_unsaved(self, cursor):
        try:
            return self.write_log[cursor]
        except KeyError:
            return self.read_saved(cursor)

    def read_saved(self, cursor):
        try:
            return self.commit_log[cursor]
        except KeyError:
            return self.read_source(cursor)

    def read_source(self, cursor):
        try:
            return self.read_log[cursor]
        except KeyError:
            self.notify()
            return log_read(self, cursor)

    def write(self, cursor):
        try:
            return self.write_log[cursor]
        except:
            return log_write(self, cursor)

    def delete(self, cursor):
        log_delete(self, cursor)

    def save_state(self, cursor, force=False):
        try:
            state = self.write_log.pop(cursor)
        except KeyError:
            ## There's nothing to write; just return.
            if force:
                self.commit_log[cursor] = self.read_saved(cursor)
                return True
            else:
                return False
        else:
            self.update_commit_log(cursor, state)
            return True

    def revert_state(self, cursor):
        self.write_log.pop(cursor, None)

    def commit_changes(self, nested):
        ## A journal is single-threaded; state can be blindly copied
        ## in.
        for (cursor, orig, state) in nested.changed():
            self.update_commit_log(cursor, state)
        nested.committed()

    def unsaved(self):
        return (
            k for (k, v) in self.write_log
            if read_saved(self, k, None) is not v
        )

    def changed(self):
        return (
            change(k, get_state(self.read_log, k), v)
            for (k, v) in self.commit_log
        )

    def read(self):
        return iter(self.read_log)

    def update_read_log(self, cursor, state):
        self.read_log[cursor] = state

    def update_write_log(self, cursor, state):
        self.write_log[cursor] = state

    def update_commit_log(self, cursor, state):
        self.commit_log[cursor] = copy_state(state)

class memory(Memory):
    JournalType = journal
    LogType = weaklog
    name = None

    def __init__(self, name='*memory*', check_read=True, check_write=True):
        self.name = name
        self.write_lock = threading.RLock()
        self.check_read = check_read
        self.check_write = check_write
        self.mem = self.LogType()

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, str(self))

    def __str__(self):
        return self.name

    def __contains__(self, cursor):
        return cursor in self.mem

    def begin(self, nested):
        pass

    def committed(self):
        pass

    def make_journal(self, name):
        return self.JournalType(name, self)

    def allocate(self, cursor, state):
        self.mem.allocate(cursor, state)

    def read_saved(self, cursor):
        return self.mem[cursor]

    def commit_changes(self, nested):
        self.with_changes(nested, nested.changed(), self.write_changes)

    def with_changes(self, trans, changed, op):
        with self.write_lock:
            self.verify_read(trans.read())
            op(trans, self.verify_write(changed))
            trans.committed()

    def verify_read(self, read):
        if self.check_read:
            verify_read(self.mem, read)

    def verify_write(self, changed):
        if self.check_write:
            return verify_write(self.mem, changed)
        else:
            return unverified_write(changed)

    def write_changes(self, nested, changed):
        for (cursor, state) in changed:
            if is_deleted(state):
                self.expire(cursor)
            else:
                self.mem[cursor] = state

    def expire(self, cursor):
        if isinstance(cursor, Cursor):
            self.mem.pop(cursor, None)
        elif callable(cursor):
            for cur in self.mem.keys():
                if cursor(cur):
                    del self.mem[cur]
        else:
            raise TypeError('expecting cursor or callable', cursor)

    def expire_all(self, cursors=None):
        if cursors:
            for cursor in cursors:
                self.expire(cursor)
        else:
            self.mem.clear()

### State

copy_state = copy.deepcopy

class sentinal(object):
    __slots__ = ('name')

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return '<%s %s>' % (type(self).__name__, self.name)

    def __copy__(self):
        return self

    def __deepcopy__(self, memo):
        return self

INSERTED = sentinal('inserted')
DELETED = sentinal('deleted')

def good_value(cursor, value, *default):
    if value is INSERTED or value is DELETED:
        if default:
            return default[0]
        else:
            raise ValueError, id(cursor)
    else:
        return value

def is_deleted(obj):
    return obj is DELETED

def is_inserted(obj):
    return obj is INSERTED

### Logged Operations

def log_read(journal, cursor):
    state = read_saved(journal.source, cursor, INSERTED)
    journal.update_read_log(cursor, state)
    return state

def log_write(journal, cursor):
    if isinstance(cursor.__id__, Static):
        raise TypeError('static cursors are not writable', cursor)
    state = copy_state(read_saved(journal, cursor))
    journal.update_write_log(cursor, state)
    return state

def log_delete(journal, cursor):
    if isinstance(cursor.__id__, Static):
        raise TypeError('static cursors are not writable', cursor)
    journal.update_write_log(cursor, DELETED)

### Operations on Logs

def get_state(log, cursor):
    return log.get(cursor, INSERTED)

def verify_read(log, read):
    conflicts = [(c, s) for (c, s) in read if log.get(c) != s]
    if conflicts:
        raise CannotCommit, conflicts

def verify_write(log, changed):
    changed, conflicts = partition_conflicts(log, changed)
    if conflicts:
        raise CannotCommit, conflicts
    return changed

def unverified_write(changed):
    return list((c, s) for (c, o, s) in changed)

def partition_conflicts(log, changed):
    good = []; bad = []
    for (cursor, orig, state) in changed:
        current = get_state(log, cursor)
        (good if current is orig else bad).append((cursor, state))
    return good, bad
<|endoftext|>


<commit_before>import sys
import os
import yaml

from twisted.python import usage
from twisted.python.util import spewer
from twisted.internet import defer, reactor

from ooni import errors, __version__

from ooni.settings import config
from ooni.director import Director
from ooni.deck import Deck, nettest_to_path
from ooni.nettest import NetTestLoader
from ooni.utils.txscapy import ScapyFactory

from ooni.utils import log, checkForRoot

class Options(usage.Options):
    synopsis = """%s [options] [path to test].py
    """ % (os.path.basename(sys.argv[0]),)

    longdesc = ("ooniprobe loads and executes a suite or a set of suites of"
                " network tests. These are loaded from modules, packages and"
                " files listed on the command line")

    optFlags = [["help", "h"],
                ["resume", "r"],
                ["no-collector", "n"],
                ["no-geoip", "g"],
                ["list", "s"],
                ["printdeck", "p"],
                ["verbose", "v"]
                ]

    optParameters = [
        ["reportfile", "o", None, "report file name"],
        ["testdeck", "i", None,
         "Specify as input a test deck: a yaml file containing the tests to run and their arguments"],
        ["collector", "c", None,
         "Address of the collector of test results. This option should not be used, but you should always use a bouncer."],
        ["bouncer", "b", 'httpo://nkvphnp3p6agi5qq.onion',
         "Address of the bouncer for test helpers. default: httpo://nkvphnp3p6agi5qq.onion"],
        ["logfile", "l", None, "log file name"],
        ["pcapfile", "O", None, "pcap file name"],
        ["configfile", "f", None,
         "Specify a path to the ooniprobe configuration file"],
        ["datadir", "d", None,
         "Specify a path to the ooniprobe data directory"],
        ["annotations", "a", None,
         "Annotate the report with a key:value[, key:value] format."]]

    compData = usage.Completions(
        extraActions=[usage.CompleteFiles(
            "*.py", descr="file | module | package | TestCase | testMethod",
            repeat=True)],)

    tracer = None

    def __init__(self):
        self['test'] = None
        usage.Options.__init__(self)

    def opt_spew(self):
        """
        Print an insanely verbose log of everything that happens.  Useful
        when debugging freezes or locks in complex code.
        """
        sys.settrace(spewer)

    def opt_version(self):
        """
        Display the ooniprobe version and exit.
        """
        print "ooniprobe version:", __version__
        sys.exit(0)

    def parseArgs(self, *args):
        if self['testdeck'] or self['list']:
            return
        try:
            self['test_file'] = args[0]
            self['subargs'] = args[1:]
        except:
            raise usage.UsageError("No test filename specified!")

def parseOptions():
    print "WARNING: running ooniprobe involves some risk that varies greatly"
    print "         from country to country. You should be aware of this when"
    print "         running the tool. Read more about this in the manpage or README."
    cmd_line_options = Options()
    if len(sys.argv) == 1:
        cmd_line_options.getUsage()
    try:
        cmd_line_options.parseOptions()
    except usage.UsageError as ue:
        print cmd_line_options.getUsage()
        raise SystemExit("%s: %s" % (sys.argv[0], ue))

    return dict(cmd_line_options)

def runWithDirector(logging=True, start_tor=True):
    """
    Instance the director, parse command line options and start an ooniprobe
    test!
    """
    global_options = parseOptions()
    config.global_options = global_options
    config.set_paths()
    config.initialize_ooni_home()
    try:
        config.read_config_file(check_incoherences=True)
    except errors.ConfigFileIncoherent:
        sys.exit(6)

    if global_options['verbose']:
        config.advanced.debug = True

    if not start_tor:
        config.advanced.start_tor = False

    if logging:
        log.start(global_options['logfile'])

    config.scapyFactory = ScapyFactory(config.advanced.interface)
    if config.privacy.includepcap:
        try:
            checkForRoot()
        except errors.InsufficientPrivileges:
            log.err("Insufficient Privileges to capture packets."
                    " See ooniprobe.conf privacy.includepcap")
            sys.exit(2)

    director = Director()
    if global_options['list']:
        print "# Installed nettests"
        for net_test_id, net_test in director.getNetTests().items():
            print "* %s (%s/%s)" % (net_test['name'],
                                    net_test['category'],
                                    net_test['id'])
            print "  %s" % net_test['description']

        sys.exit(0)

    elif global_options['printdeck']:
        del global_options['printdeck']
        print "# Copy and paste the lines below into a test deck to run the specified test with the specified arguments"
        print yaml.safe_dump([{'options': global_options}]).strip()

        sys.exit(0)

    if global_options.get('annotations') is not None:
        annotations = {}
        for annotation in global_options["annotations"].split(","):
            pair = annotation.split(":")
            if len(pair) == 2:
                key = pair[0].strip()
                value = pair[1].strip()
                annotations[key] = value
            else:
                log.err("Invalid annotation: %s" % annotation)
                sys.exit(1)
        global_options["annotations"] = annotations

    # XXX: This should mean no bouncer either!
    if global_options['no-collector']:
        log.msg("Not reporting using a collector")
        global_options['collector'] = None
        global_options['bouncer'] = None

    deck = Deck()
    deck.bouncer = global_options['bouncer']
    conditions_to_start_tor = [False]
    if global_options['bouncer']:
        conditions_to_start_tor.append(True)
    if global_options['collector']:
        conditions_to_start_tor.append(True)

    try:
        if global_options['testdeck']:
            deck.loadDeck(global_options['testdeck'])
        else:
            log.debug("No test deck detected")
            test_file = nettest_to_path(global_options['test_file'], True)
            net_test_loader = NetTestLoader(global_options['subargs'],
                                            test_file=test_file)
            deck.insert(net_test_loader)
    except errors.MissingRequiredOption as option_name:
        log.err('Missing required option: "%s"' % option_name)
        incomplete_net_test_loader = option_name.net_test_loader
        print incomplete_net_test_loader.usageOptions().getUsage()
        sys.exit(2)
    except errors.NetTestNotFound as path:
        log.err('Requested NetTest file not found (%s)' % path)
        sys.exit(3)
    except errors.OONIUsageError as e:
        log.err(e)
        print e.net_test_loader.usageOptions().getUsage()
        sys.exit(4)
    except Exception as e:
        log.err(e)
        sys.exit(5)

    conditions_to_start_tor.append(deck.requiresTor)
    start_tor = any(conditions_to_start_tor)
    d = director.start(start_tor=start_tor)

    def setup_nettest(_):
        try:
            return deck.setup()
        except errors.UnableToLoadDeckInput as error:
            return defer.failure.Failure(error)

    def director_startup_handled_failures(failure):
        log.err("Could not start the director")
        failure.trap(errors.TorNotRunning,
                     errors.InvalidOONIBCollectorAddress,
                     errors.UnableToLoadDeckInput,
                     errors.CouldNotFindTestHelper,
                     errors.CouldNotFindTestCollector,
                     errors.ProbeIPUnknown,
                     errors.InvalidInputFile,
                     errors.ConfigFileIncoherent)

        if isinstance(failure.value, errors.TorNotRunning):
            log.err("Tor does not appear to be running")
            log.err("Reporting with the collector %s is not possible" %
                    global_options['collector'])
            log.msg(
                "Try with a different collector or disable collector reporting with -n")

        elif isinstance(failure.value, errors.InvalidOONIBCollectorAddress):
            log.err("Invalid format for oonib collector address.")
            log.msg(
                "Should be in the format http://<collector_address>:<port>")
            log.msg("for example: ooniprobe -c httpo://nkvphnp3p6agi5qq.onion")

        elif isinstance(failure.value, errors.UnableToLoadDeckInput):
            log.err("Unable to fetch the required inputs for the test deck.")
            log.msg(
                "Please file a ticket on our issue tracker: https://github.com/thetorproject/ooni-probe/issues")

        elif isinstance(failure.value, errors.CouldNotFindTestHelper):
            log.err("Unable to obtain the required test helpers.")
            log.msg(
                "Try with a different bouncer or check that Tor is running properly.")

        elif isinstance(failure.value, errors.CouldNotFindTestCollector):
            log.err("Could not find a valid collector.")
            log.msg(
                "Try with a different bouncer, specify a collector with -c or disable reporting to a collector with -n.")

        elif isinstance(failure.value, errors.ProbeIPUnknown):
            log.err("Failed to lookup probe IP address.")
            log.msg("Check your internet connection.")

        elif isinstance(failure.value, errors.InvalidInputFile):
            log.err("Invalid input file \"%s\"" % failure.value)

        elif isinstance(failure.value, errors.ConfigFileIncoherent):
            log.err("Incoherent config file")

        if config.advanced.debug:
            log.exception(failure)

    def director_startup_other_failures(failure):
        log.err("An unhandled exception occurred while starting the director!")
        log.exception(failure)

    # Wait until director has started up (including bootstrapping Tor)
    # before adding tests
    def post_director_start(_):
        for net_test_loader in deck.netTestLoaders:
            # Decks can specify different collectors
            # for each net test, so that each NetTest
            # may be paired with a test_helper and its collector
            # However, a user can override this behavior by
            # specifying a collector from the command-line (-c).
            # If a collector is not specified in the deck, or the
            # deck is a singleton, the default collector set in
            # ooniprobe.conf will be used

            collector = None
            if not global_options['no-collector']:
                if global_options['collector']:
                    collector = global_options['collector']
                elif 'collector' in config.reports \
                        and config.reports['collector']:
                    collector = config.reports['collector']
                elif net_test_loader.collector:
                    collector = net_test_loader.collector

            if collector and collector.startswith('httpo:') \
                    and (not (config.tor_state or config.tor.socks_port)):
                raise errors.TorNotRunning

            test_details = net_test_loader.testDetails
            test_details['annotations'] = global_options['annotations']

            director.startNetTest(net_test_loader,
                                  global_options['reportfile'],
                                  collector)
        return director.allTestsDone

    def start():
        d.addCallback(setup_nettest)
        d.addCallback(post_director_start)
        d.addErrback(director_startup_handled_failures)
        d.addErrback(director_startup_other_failures)
        return d

    return start()
<commit_msg>Refactor the condition checking to not use an array<commit_after>import sys
import os
import yaml

from twisted.python import usage
from twisted.python.util import spewer
from twisted.internet import defer, reactor

from ooni import errors, __version__

from ooni.settings import config
from ooni.director import Director
from ooni.deck import Deck, nettest_to_path
from ooni.nettest import NetTestLoader
from ooni.utils.txscapy import ScapyFactory

from ooni.utils import log, checkForRoot

class Options(usage.Options):
    synopsis = """%s [options] [path to test].py
    """ % (os.path.basename(sys.argv[0]),)

    longdesc = ("ooniprobe loads and executes a suite or a set of suites of"
                " network tests. These are loaded from modules, packages and"
                " files listed on the command line")

    optFlags = [["help", "h"],
                ["resume", "r"],
                ["no-collector", "n"],
                ["no-geoip", "g"],
                ["list", "s"],
                ["printdeck", "p"],
                ["verbose", "v"]
                ]

    optParameters = [
        ["reportfile", "o", None, "report file name"],
        ["testdeck", "i", None,
         "Specify as input a test deck: a yaml file containing the tests to run and their arguments"],
        ["collector", "c", None,
         "Address of the collector of test results. This option should not be used, but you should always use a bouncer."],
        ["bouncer", "b", 'httpo://nkvphnp3p6agi5qq.onion',
         "Address of the bouncer for test helpers. default: httpo://nkvphnp3p6agi5qq.onion"],
        ["logfile", "l", None, "log file name"],
        ["pcapfile", "O", None, "pcap file name"],
        ["configfile", "f", None,
         "Specify a path to the ooniprobe configuration file"],
        ["datadir", "d", None,
         "Specify a path to the ooniprobe data directory"],
        ["annotations", "a", None,
         "Annotate the report with a key:value[, key:value] format."]]

    compData = usage.Completions(
        extraActions=[usage.CompleteFiles(
            "*.py", descr="file | module | package | TestCase | testMethod",
            repeat=True)],)

    tracer = None

    def __init__(self):
        self['test'] = None
        usage.Options.__init__(self)

    def opt_spew(self):
        """
        Print an insanely verbose log of everything that happens.  Useful
        when debugging freezes or locks in complex code.
        """
        sys.settrace(spewer)

    def opt_version(self):
        """
        Display the ooniprobe version and exit.
        """
        print "ooniprobe version:", __version__
        sys.exit(0)

    def parseArgs(self, *args):
        if self['testdeck'] or self['list']:
            return
        try:
            self['test_file'] = args[0]
            self['subargs'] = args[1:]
        except:
            raise usage.UsageError("No test filename specified!")

def parseOptions():
    print "WARNING: running ooniprobe involves some risk that varies greatly"
    print "         from country to country. You should be aware of this when"
    print "         running the tool. Read more about this in the manpage or README."
    cmd_line_options = Options()
    if len(sys.argv) == 1:
        cmd_line_options.getUsage()
    try:
        cmd_line_options.parseOptions()
    except usage.UsageError as ue:
        print cmd_line_options.getUsage()
        raise SystemExit("%s: %s" % (sys.argv[0], ue))

    return dict(cmd_line_options)

def runWithDirector(logging=True, start_tor=True):
    """
    Instance the director, parse command line options and start an ooniprobe
    test!
    """
    global_options = parseOptions()
    config.global_options = global_options
    config.set_paths()
    config.initialize_ooni_home()
    try:
        config.read_config_file(check_incoherences=True)
    except errors.ConfigFileIncoherent:
        sys.exit(6)

    if global_options['verbose']:
        config.advanced.debug = True

    if not start_tor:
        config.advanced.start_tor = False

    if logging:
        log.start(global_options['logfile'])

    config.scapyFactory = ScapyFactory(config.advanced.interface)
    if config.privacy.includepcap:
        try:
            checkForRoot()
        except errors.InsufficientPrivileges:
            log.err("Insufficient Privileges to capture packets."
                    " See ooniprobe.conf privacy.includepcap")
            sys.exit(2)

    director = Director()
    if global_options['list']:
        print "# Installed nettests"
        for net_test_id, net_test in director.getNetTests().items():
            print "* %s (%s/%s)" % (net_test['name'],
                                    net_test['category'],
                                    net_test['id'])
            print "  %s" % net_test['description']

        sys.exit(0)

    elif global_options['printdeck']:
        del global_options['printdeck']
        print "# Copy and paste the lines below into a test deck to run the specified test with the specified arguments"
        print yaml.safe_dump([{'options': global_options}]).strip()

        sys.exit(0)

    if global_options.get('annotations') is not None:
        annotations = {}
        for annotation in global_options["annotations"].split(","):
            pair = annotation.split(":")
            if len(pair) == 2:
                key = pair[0].strip()
                value = pair[1].strip()
                annotations[key] = value
            else:
                log.err("Invalid annotation: %s" % annotation)
                sys.exit(1)
        global_options["annotations"] = annotations

    # XXX: This should mean no bouncer either!
    if global_options['no-collector']:
        log.msg("Not reporting using a collector")
        global_options['collector'] = None
        global_options['bouncer'] = None

    deck = Deck()
    deck.bouncer = global_options['bouncer']
    start_tor = False
    if global_options['bouncer']:
        start_tor |= True
    if global_options['collector']:
        start_tor |= True

    try:
        if global_options['testdeck']:
            deck.loadDeck(global_options['testdeck'])
        else:
            log.debug("No test deck detected")
            test_file = nettest_to_path(global_options['test_file'], True)
            net_test_loader = NetTestLoader(global_options['subargs'],
                                            test_file=test_file)
            deck.insert(net_test_loader)
    except errors.MissingRequiredOption as option_name:
        log.err('Missing required option: "%s"' % option_name)
        incomplete_net_test_loader = option_name.net_test_loader
        print incomplete_net_test_loader.usageOptions().getUsage()
        sys.exit(2)
    except errors.NetTestNotFound as path:
        log.err('Requested NetTest file not found (%s)' % path)
        sys.exit(3)
    except errors.OONIUsageError as e:
        log.err(e)
        print e.net_test_loader.usageOptions().getUsage()
        sys.exit(4)
    except Exception as e:
        log.err(e)
        sys.exit(5)

    start_tor |= deck.requiresTor
    d = director.start(start_tor=start_tor)

    def setup_nettest(_):
        try:
            return deck.setup()
        except errors.UnableToLoadDeckInput as error:
            return defer.failure.Failure(error)

    def director_startup_handled_failures(failure):
        log.err("Could not start the director")
        failure.trap(errors.TorNotRunning,
                     errors.InvalidOONIBCollectorAddress,
                     errors.UnableToLoadDeckInput,
                     errors.CouldNotFindTestHelper,
                     errors.CouldNotFindTestCollector,
                     errors.ProbeIPUnknown,
                     errors.InvalidInputFile,
                     errors.ConfigFileIncoherent)

        if isinstance(failure.value, errors.TorNotRunning):
            log.err("Tor does not appear to be running")
            log.err("Reporting with the collector %s is not possible" %
                    global_options['collector'])
            log.msg(
                "Try with a different collector or disable collector reporting with -n")

        elif isinstance(failure.value, errors.InvalidOONIBCollectorAddress):
            log.err("Invalid format for oonib collector address.")
            log.msg(
                "Should be in the format http://<collector_address>:<port>")
            log.msg("for example: ooniprobe -c httpo://nkvphnp3p6agi5qq.onion")

        elif isinstance(failure.value, errors.UnableToLoadDeckInput):
            log.err("Unable to fetch the required inputs for the test deck.")
            log.msg(
                "Please file a ticket on our issue tracker: https://github.com/thetorproject/ooni-probe/issues")

        elif isinstance(failure.value, errors.CouldNotFindTestHelper):
            log.err("Unable to obtain the required test helpers.")
            log.msg(
                "Try with a different bouncer or check that Tor is running properly.")

        elif isinstance(failure.value, errors.CouldNotFindTestCollector):
            log.err("Could not find a valid collector.")
            log.msg(
                "Try with a different bouncer, specify a collector with -c or disable reporting to a collector with -n.")

        elif isinstance(failure.value, errors.ProbeIPUnknown):
            log.err("Failed to lookup probe IP address.")
            log.msg("Check your internet connection.")

        elif isinstance(failure.value, errors.InvalidInputFile):
            log.err("Invalid input file \"%s\"" % failure.value)

        elif isinstance(failure.value, errors.ConfigFileIncoherent):
            log.err("Incoherent config file")

        if config.advanced.debug:
            log.exception(failure)

    def director_startup_other_failures(failure):
        log.err("An unhandled exception occurred while starting the director!")
        log.exception(failure)

    # Wait until director has started up (including bootstrapping Tor)
    # before adding tests
    def post_director_start(_):
        for net_test_loader in deck.netTestLoaders:
            # Decks can specify different collectors
            # for each net test, so that each NetTest
            # may be paired with a test_helper and its collector
            # However, a user can override this behavior by
            # specifying a collector from the command-line (-c).
            # If a collector is not specified in the deck, or the
            # deck is a singleton, the default collector set in
            # ooniprobe.conf will be used

            collector = None
            if not global_options['no-collector']:
                if global_options['collector']:
                    collector = global_options['collector']
                elif 'collector' in config.reports \
                        and config.reports['collector']:
                    collector = config.reports['collector']
                elif net_test_loader.collector:
                    collector = net_test_loader.collector

            if collector and collector.startswith('httpo:') \
                    and (not (config.tor_state or config.tor.socks_port)):
                raise errors.TorNotRunning

            test_details = net_test_loader.testDetails
            test_details['annotations'] = global_options['annotations']

            director.startNetTest(net_test_loader,
                                  global_options['reportfile'],
                                  collector)
        return director.allTestsDone

    def start():
        d.addCallback(setup_nettest)
        d.addCallback(post_director_start)
        d.addErrback(director_startup_handled_failures)
        d.addErrback(director_startup_other_failures)
        return d

    return start()
<|endoftext|>


<commit_before>"""
Equivalencies between different kinds of units

"""

# -----------------------------------------------------------------------------
# Copyright (c) 2018, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# -----------------------------------------------------------------------------

from collections import OrderedDict

from unyt.dimensions import (
    temperature,
    mass,
    energy,
    length,
    rate,
    velocity,
    dimensionless,
    density,
    number_density,
    flux,
    current_cgs,
    current_mks,
    charge_cgs,
    charge_mks,
    magnetic_field_cgs,
    magnetic_field_mks,
    electric_potential_cgs,
    electric_potential_mks,
    electric_field_cgs,
    electric_field_mks,
    resistance_cgs,
    resistance_mks
)

from unyt._physical_ratios import speed_of_light_cm_per_s
from six import add_metaclass
import numpy as np

equivalence_registry = OrderedDict()

class _RegisteredEquivalence(type):
    def __init__(cls, name, b, d):
        type.__init__(cls, name, b, d)
        if hasattr(cls, "type_name"):
            equivalence_registry[cls.type_name] = cls
        if hasattr(cls, "alternate_names"):
            for name in cls.alternate_names:
                equivalence_registry[name] = cls

@add_metaclass(_RegisteredEquivalence)
class Equivalence(object):
    one_way = False

class NumberDensityEquivalence(Equivalence):
    """Equivalence between mass and number density, given a mean molecular weight.

    Given a number density :math:`n`, the mass density :math:`\\rho` is:

    .. math::

      \\rho = \\mu m_{\\rm H} n

    And similarly

    .. math::

      n = \\rho (\\mu m_{\\rm H})^{-1}

    Parameters
    ----------
    mu : float
      The mean molecular weight. Defaults to 0.6 whcih is valid for fully
      ionized gas with primordial composition.

    Example
    -------
    >>> print(NumberDensityEquivalence())
    number density: density <-> number density
    >>> from unyt import Msun, pc
    >>> rho = Msun/pc**3
    >>> rho.to_equivalent('cm**-3', 'number_density', mu=1.4)
    unyt_quantity(28.88289965, 'cm**(-3)')
    """
    type_name = "number_density"
    _dims = (density, number_density,)

    def _convert(self, x, new_dims, mu=0.6):
        from unyt import physical_constants as pc
        if new_dims == number_density:
            return x/(mu*pc.mh)
        elif new_dims == density:
            return x*mu*pc.mh

    def __str__(self):
        return "number density: density <-> number density"

class ThermalEquivalence(Equivalence):
    """Equivalence between temperature and energy via the Boltzmann constant

    Given a temperature :math:`T` in an absolute scale (e.g. Kelvin or
    Rankine), the equivalent thermal energy :math:`E` for that temperature is
    given by:

    .. math::

      E = k_B T

    And

    .. math::

      T = E/k_B

    Where :math:`k_B` is Boltzmann's constant.

    Example
    -------
    >>> print(ThermalEquivalence())
    thermal: temperature <-> energy
    >>> from unyt import Kelvin
    >>> temp = 1e6*Kelvin
    >>> temp.to_equivalent('keV', 'thermal')
    unyt_quantity(0.08617332, 'keV')
    """
    type_name = "thermal"
    _dims = (temperature, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            return pc.kboltz*x
        elif new_dims == temperature:
            return x/pc.kboltz

    def __str__(self):
        return "thermal: temperature <-> energy"

class MassEnergyEquivalence(Equivalence):
    """Equivalence between mass and energy in special relativity

    Given a body with mass :math:`m`, the self-energy :math:`E` of that mass is
    given by

    .. math::

      E = m c^2

    where :math:`c` is the speed of light.

    Example
    -------
    >>> print(MassEnergyEquivalence())
    mass_energy: mass <-> energy
    >>> from unyt import g
    >>> g.to_equivalent('J', 'mass_energy')
    unyt_quantity(8.98755179e+13, 'J')

    """
    type_name = "mass_energy"
    _dims = (mass, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            return x*pc.clight*pc.clight
        elif new_dims == mass:
            return x/(pc.clight*pc.clight)

    def __str__(self):
        return "mass_energy: mass <-> energy"

class SpectralEquivalence(Equivalence):
    """Equivalence between wavelength, frequency, and energy of a photon.

    Given a photon with wavelength :math:`\\lambda`, frequency :math:`\\nu`
    and Energy :math:`E`, these quantities are related by the following
    forumlae:

    .. math::

      E = h \\nu = h c / \\lambda

    where :math:`h` is Planck's constant and :math:`c` is the speed of light.

    Example
    ------
    >>> print(SpectralEquivalence())
    spectral: length <-> frequency <-> energy
    >>> from unyt import angstrom, km
    >>> angstrom.to_equivalent('keV', 'spectral')
    unyt_quantity(12.39841932, 'keV')
    >>> km.to_equivalent('MHz', 'spectral')
    unyt_quantity(0.29979246, 'MHz')
    """
    type_name = "spectral"
    _dims = (length, rate, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            if x.units.dimensions == length:
                nu = pc.clight/x
            elif x.units.dimensions == rate:
                nu = x
            return pc.hcgs*nu
        elif new_dims == length:
            if x.units.dimensions == rate:
                return pc.clight/x
            elif x.units.dimensions == energy:
                return pc.hcgs*pc.clight/x
        elif new_dims == rate:
            if x.units.dimensions == length:
                return pc.clight/x
            elif x.units.dimensions == energy:
                return x/pc.hcgs

    def __str__(self):
        return "spectral: length <-> frequency <-> energy"

class SoundSpeedEquivalence(Equivalence):
    """Equivalence between the sound speed, temperature, and thermal energy of
    an ideal gas

    For an ideal gas with sound speed :math:`c_s`, temperature :math:`T`, and
    thermal energy :math:`E`, the following equalities will hold:

    .. math::

      c_s = \\sqrt{\\frac{\\gamma k_B T}{\\mu m_{\\rm H}}

    and

    .. math::

      E = c_s^2 \\mu m_{\\rm H} / \\gammma = k_B T

    where :math:`k_B` is Boltzmann's constant, :math:`\\mu` is the mean
    molecular weight of the gas, and :math:`\\gamma` is the ratio of specific
    heats.

    Parameters
    ----------
    gamma : float
       The ratio of specific heats. Defaults to 5/3, which is correct for
       monatomic species.
    mu : float
       The mean molecular weight. Defaults to 0.6, which is valid for fully
       ionized gas with primordial composition.

    Example
    -------
    >>> print(SoundSpeedEquivalence())
    sound_speed (ideal gas): velocity <-> temperature <-> energy
    >>> from unyt import Kelvin
    >>> hot = 1e6*Kelvin
    >>> hot.to_equivalent('km/s', 'sound_speed')
    unyt_quantity(151.37249927, 'km/s')
    >>> hot.to_equivalent('keV', 'sound_speed')
    unyt_quantity(0.08617332, 'keV')
    """
    type_name = "sound_speed"
    _dims = (velocity, temperature, energy,)

    def _convert(self, x, new_dims, mu=0.6, gamma=5./3.):
        from unyt import physical_constants as pc
        if new_dims == velocity:
            if x.units.dimensions == temperature:
                kT = pc.kboltz*x
            elif x.units.dimensions == energy:
                kT = x
            return np.sqrt(gamma*kT/(mu*pc.mh))
        elif new_dims == temperature:
            if x.units.dimensions == velocity:
                kT = x*x*mu*pc.mh/gamma
                return kT/pc.kboltz
            else:
                return x/pc.kboltz
        else:
            if x.units.dimensions == velocity:
                return np.sqrt(gamma*x/(mu*pc.mh))
            else:
                return x*pc.kboltz

    def __str__(self):
        return "sound_speed (ideal gas): velocity <-> temperature <-> energy"

class LorentzEquivalence(Equivalence):
    """Equivalence between velocity and the Lorentz gamma factor.

    For a body with velocity :math:`v`, the Lorentz gamma factor,
    :math:`\\gamma` is

    .. math::

      \\gamma = \\frac{1}{\sqrt(1 - v^2/c^2}}

    where :math:`c` is the speed of light.

    Example
    -------
    >>> from unyt import c
    >>> v = 0.99*c
    >>> print(v.to_equivalent('', 'lorentz'))
    7.088812050083354 dimensionless
    """
    type_name = "lorentz"
    _dims = (dimensionless, velocity,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == dimensionless:
            beta = x.in_cgs()/pc.clight
            return 1./np.sqrt(1.-beta**2)
        elif new_dims == velocity:
            return pc.clight*np.sqrt(1.-1./(x*x))

    def __str__(self):
        return "lorentz: velocity <-> dimensionless"

class SchwarzschildEquivalence(Equivalence):
    """Equivalence between the mass and radius of a Schwarzschild black hole

    A Schwarzschild black hole of mass :math:`M` has radius :math:`R`

    .. math::

      R = \\frac{2 G M}{c^2}

    and similarly

    .. math::

      M = \\frac{R c^2}{2 G}

    where :math:`G` is Newton's gravitational constant and :math:`c` is the
    speed of light.

    Example
    -------
    >>> from unyt import Msun
    >>> Msun.to_equivalent('km', 'schwarzschild')
    unyt_quantity(2.95305543, 'km')
    """
    type_name = "schwarzschild"
    _dims = (mass, length,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == length:
            return 2.*pc.G*x/(pc.clight*pc.clight)
        elif new_dims == mass:
            return 0.5*x*pc.clight*pc.clight/pc.G

    def __str__(self):
        return "schwarzschild: mass <-> length"

class ComptonEquivalence(Equivalence):
    type_name = "compton"
    _dims = (mass, length,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        return pc.hcgs/(x*pc.clight)

    def __str__(self):
        return "compton: mass <-> length"

class EffectiveTemperature(Equivalence):
    type_name = "effective_temperature"
    _dims = (flux, temperature,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == flux:
            return pc.stefan_boltzmann_constant_cgs*x**4
        elif new_dims == temperature:
            return (x/pc.stefan_boltzmann_constant_cgs)**0.25

    def __str__(self):
        return "effective_temperature: flux <-> temperature"

em_conversions = {
    charge_mks: ("esu", 0.1*speed_of_light_cm_per_s),
    magnetic_field_mks: ("gauss", 1.0e4),
    current_mks: ("statA", 0.1*speed_of_light_cm_per_s),
    electric_potential_mks: ("statV", 1.0e-8*speed_of_light_cm_per_s),
    resistance_mks: ("statohm", 1.0e9/(speed_of_light_cm_per_s**2)),
    charge_cgs: ("C", 10.0/speed_of_light_cm_per_s),
    magnetic_field_cgs: ("T", 1.0e-4),
    current_cgs: ("A", 10.0/speed_of_light_cm_per_s),
    electric_potential_cgs: ("V", 1.0e8/speed_of_light_cm_per_s),
    resistance_cgs: ("ohm", speed_of_light_cm_per_s**2*1.0e-9),
}

class ElectromagneticSI(Equivalence):
    type_name = "SI"
    alternate_names = ["si", "MKS", "mks"]
    one_way = True
    _dims = (current_cgs, charge_cgs, magnetic_field_cgs,
             electric_field_cgs, electric_potential_cgs,
             resistance_cgs)

    def _convert(self, x, new_dims):
        old_dims = x.units.dimensions
        new_units, convert_factor = em_conversions[old_dims]
        return x.in_cgs().v*convert_factor, new_units

    def __str__(self):
        return "SI: EM CGS unit -> EM SI unit"

class ElectromagneticCGS(Equivalence):
    type_name = "CGS"
    alternate_names = ["cgs"]
    one_way = True
    _dims = (current_mks, charge_mks, magnetic_field_mks,
             electric_field_mks, electric_potential_mks,
             resistance_mks)

    def _convert(self, x, new_dims):
        old_dims = x.units.dimensions
        new_units, convert_factor = em_conversions[old_dims]
        return x.in_mks().v*convert_factor, new_units

    def __str__(self):
        return "CGS: EM SI unit -> EM CGS unit"
<commit_msg>Fix LaTeX formatting issues<commit_after>"""
Equivalencies between different kinds of units

"""

# -----------------------------------------------------------------------------
# Copyright (c) 2018, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# -----------------------------------------------------------------------------

from collections import OrderedDict

from unyt.dimensions import (
    temperature,
    mass,
    energy,
    length,
    rate,
    velocity,
    dimensionless,
    density,
    number_density,
    flux,
    current_cgs,
    current_mks,
    charge_cgs,
    charge_mks,
    magnetic_field_cgs,
    magnetic_field_mks,
    electric_potential_cgs,
    electric_potential_mks,
    electric_field_cgs,
    electric_field_mks,
    resistance_cgs,
    resistance_mks
)

from unyt._physical_ratios import speed_of_light_cm_per_s
from six import add_metaclass
import numpy as np

equivalence_registry = OrderedDict()

class _RegisteredEquivalence(type):
    def __init__(cls, name, b, d):
        type.__init__(cls, name, b, d)
        if hasattr(cls, "type_name"):
            equivalence_registry[cls.type_name] = cls
        if hasattr(cls, "alternate_names"):
            for name in cls.alternate_names:
                equivalence_registry[name] = cls

@add_metaclass(_RegisteredEquivalence)
class Equivalence(object):
    one_way = False

class NumberDensityEquivalence(Equivalence):
    """Equivalence between mass and number density, given a mean molecular weight.

    Given a number density :math:`n`, the mass density :math:`\\rho` is:

    .. math::

      \\rho = \\mu m_{\\rm H} n

    And similarly

    .. math::

      n = \\rho (\\mu m_{\\rm H})^{-1}

    Parameters
    ----------
    mu : float
      The mean molecular weight. Defaults to 0.6 whcih is valid for fully
      ionized gas with primordial composition.

    Example
    -------
    >>> print(NumberDensityEquivalence())
    number density: density <-> number density
    >>> from unyt import Msun, pc
    >>> rho = Msun/pc**3
    >>> rho.to_equivalent('cm**-3', 'number_density', mu=1.4)
    unyt_quantity(28.88289965, 'cm**(-3)')
    """
    type_name = "number_density"
    _dims = (density, number_density,)

    def _convert(self, x, new_dims, mu=0.6):
        from unyt import physical_constants as pc
        if new_dims == number_density:
            return x/(mu*pc.mh)
        elif new_dims == density:
            return x*mu*pc.mh

    def __str__(self):
        return "number density: density <-> number density"

class ThermalEquivalence(Equivalence):
    """Equivalence between temperature and energy via the Boltzmann constant

    Given a temperature :math:`T` in an absolute scale (e.g. Kelvin or
    Rankine), the equivalent thermal energy :math:`E` for that temperature is
    given by:

    .. math::

      E = k_B T

    And

    .. math::

      T = E/k_B

    Where :math:`k_B` is Boltzmann's constant.

    Example
    -------
    >>> print(ThermalEquivalence())
    thermal: temperature <-> energy
    >>> from unyt import Kelvin
    >>> temp = 1e6*Kelvin
    >>> temp.to_equivalent('keV', 'thermal')
    unyt_quantity(0.08617332, 'keV')
    """
    type_name = "thermal"
    _dims = (temperature, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            return pc.kboltz*x
        elif new_dims == temperature:
            return x/pc.kboltz

    def __str__(self):
        return "thermal: temperature <-> energy"

class MassEnergyEquivalence(Equivalence):
    """Equivalence between mass and energy in special relativity

    Given a body with mass :math:`m`, the self-energy :math:`E` of that mass is
    given by

    .. math::

      E = m c^2

    where :math:`c` is the speed of light.

    Example
    -------
    >>> print(MassEnergyEquivalence())
    mass_energy: mass <-> energy
    >>> from unyt import g
    >>> g.to_equivalent('J', 'mass_energy')
    unyt_quantity(8.98755179e+13, 'J')

    """
    type_name = "mass_energy"
    _dims = (mass, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            return x*pc.clight*pc.clight
        elif new_dims == mass:
            return x/(pc.clight*pc.clight)

    def __str__(self):
        return "mass_energy: mass <-> energy"

class SpectralEquivalence(Equivalence):
    """Equivalence between wavelength, frequency, and energy of a photon.

    Given a photon with wavelength :math:`\\lambda`, frequency :math:`\\nu`
    and Energy :math:`E`, these quantities are related by the following
    forumlae:

    .. math::

      E = h \\nu = h c / \\lambda

    where :math:`h` is Planck's constant and :math:`c` is the speed of light.

    Example
    ------
    >>> print(SpectralEquivalence())
    spectral: length <-> frequency <-> energy
    >>> from unyt import angstrom, km
    >>> angstrom.to_equivalent('keV', 'spectral')
    unyt_quantity(12.39841932, 'keV')
    >>> km.to_equivalent('MHz', 'spectral')
    unyt_quantity(0.29979246, 'MHz')
    """
    type_name = "spectral"
    _dims = (length, rate, energy,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == energy:
            if x.units.dimensions == length:
                nu = pc.clight/x
            elif x.units.dimensions == rate:
                nu = x
            return pc.hcgs*nu
        elif new_dims == length:
            if x.units.dimensions == rate:
                return pc.clight/x
            elif x.units.dimensions == energy:
                return pc.hcgs*pc.clight/x
        elif new_dims == rate:
            if x.units.dimensions == length:
                return pc.clight/x
            elif x.units.dimensions == energy:
                return x/pc.hcgs

    def __str__(self):
        return "spectral: length <-> frequency <-> energy"

class SoundSpeedEquivalence(Equivalence):
    """Equivalence between the sound speed, temperature, and thermal energy of
    an ideal gas

    For an ideal gas with sound speed :math:`c_s`, temperature :math:`T`, and
    thermal energy :math:`E`, the following equalities will hold:

    .. math::

      c_s = \\sqrt{\\frac{\\gamma k_B T}{\\mu m_{\\rm H}}}

    and

    .. math::

      E = c_s^2 \\mu m_{\\rm H} / \\gamma = k_B T

    where :math:`k_B` is Boltzmann's constant, :math:`\\mu` is the mean
    molecular weight of the gas, and :math:`\\gamma` is the ratio of specific
    heats.

    Parameters
    ----------
    gamma : float
       The ratio of specific heats. Defaults to 5/3, which is correct for
       monatomic species.
    mu : float
       The mean molecular weight. Defaults to 0.6, which is valid for fully
       ionized gas with primordial composition.

    Example
    -------
    >>> print(SoundSpeedEquivalence())
    sound_speed (ideal gas): velocity <-> temperature <-> energy
    >>> from unyt import Kelvin
    >>> hot = 1e6*Kelvin
    >>> hot.to_equivalent('km/s', 'sound_speed')
    unyt_quantity(151.37249927, 'km/s')
    >>> hot.to_equivalent('keV', 'sound_speed')
    unyt_quantity(0.08617332, 'keV')
    """
    type_name = "sound_speed"
    _dims = (velocity, temperature, energy,)

    def _convert(self, x, new_dims, mu=0.6, gamma=5./3.):
        from unyt import physical_constants as pc
        if new_dims == velocity:
            if x.units.dimensions == temperature:
                kT = pc.kboltz*x
            elif x.units.dimensions == energy:
                kT = x
            return np.sqrt(gamma*kT/(mu*pc.mh))
        elif new_dims == temperature:
            if x.units.dimensions == velocity:
                kT = x*x*mu*pc.mh/gamma
                return kT/pc.kboltz
            else:
                return x/pc.kboltz
        else:
            if x.units.dimensions == velocity:
                return np.sqrt(gamma*x/(mu*pc.mh))
            else:
                return x*pc.kboltz

    def __str__(self):
        return "sound_speed (ideal gas): velocity <-> temperature <-> energy"

class LorentzEquivalence(Equivalence):
    """Equivalence between velocity and the Lorentz gamma factor.

    For a body with velocity :math:`v`, the Lorentz gamma factor,
    :math:`\\gamma` is

    .. math::

      \\gamma = \\frac{1}{\\sqrt{1 - v^2/c^2}}

    ans similarly

    .. math::

      v = \\frac{c}{\\sqrt{1 - \\gamma^2}}

    where :math:`c` is the speed of light.

    Example
    -------
    >>> from unyt import c
    >>> v = 0.99*c
    >>> print(v.to_equivalent('', 'lorentz'))
    7.088812050083354 dimensionless
    """
    type_name = "lorentz"
    _dims = (dimensionless, velocity,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == dimensionless:
            beta = x.in_cgs()/pc.clight
            return 1./np.sqrt(1.-beta**2)
        elif new_dims == velocity:
            return pc.clight*np.sqrt(1.-1./(x*x))

    def __str__(self):
        return "lorentz: velocity <-> dimensionless"

class SchwarzschildEquivalence(Equivalence):
    """Equivalence between the mass and radius of a Schwarzschild black hole

    A Schwarzschild black hole of mass :math:`M` has radius :math:`R`

    .. math::

      R = \\frac{2 G M}{c^2}

    and similarly

    .. math::

      M = \\frac{R c^2}{2 G}

    where :math:`G` is Newton's gravitational constant and :math:`c` is the
    speed of light.

    Example
    -------
    >>> from unyt import Msun
    >>> Msun.to_equivalent('km', 'schwarzschild')
    unyt_quantity(2.95305543, 'km')
    """
    type_name = "schwarzschild"
    _dims = (mass, length,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == length:
            return 2.*pc.G*x/(pc.clight*pc.clight)
        elif new_dims == mass:
            return 0.5*x*pc.clight*pc.clight/pc.G

    def __str__(self):
        return "schwarzschild: mass <-> length"

class ComptonEquivalence(Equivalence):
    type_name = "compton"
    _dims = (mass, length,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        return pc.hcgs/(x*pc.clight)

    def __str__(self):
        return "compton: mass <-> length"

class EffectiveTemperature(Equivalence):
    type_name = "effective_temperature"
    _dims = (flux, temperature,)

    def _convert(self, x, new_dims):
        from unyt import physical_constants as pc
        if new_dims == flux:
            return pc.stefan_boltzmann_constant_cgs*x**4
        elif new_dims == temperature:
            return (x/pc.stefan_boltzmann_constant_cgs)**0.25

    def __str__(self):
        return "effective_temperature: flux <-> temperature"

em_conversions = {
    charge_mks: ("esu", 0.1*speed_of_light_cm_per_s),
    magnetic_field_mks: ("gauss", 1.0e4),
    current_mks: ("statA", 0.1*speed_of_light_cm_per_s),
    electric_potential_mks: ("statV", 1.0e-8*speed_of_light_cm_per_s),
    resistance_mks: ("statohm", 1.0e9/(speed_of_light_cm_per_s**2)),
    charge_cgs: ("C", 10.0/speed_of_light_cm_per_s),
    magnetic_field_cgs: ("T", 1.0e-4),
    current_cgs: ("A", 10.0/speed_of_light_cm_per_s),
    electric_potential_cgs: ("V", 1.0e8/speed_of_light_cm_per_s),
    resistance_cgs: ("ohm", speed_of_light_cm_per_s**2*1.0e-9),
}

class ElectromagneticSI(Equivalence):
    type_name = "SI"
    alternate_names = ["si", "MKS", "mks"]
    one_way = True
    _dims = (current_cgs, charge_cgs, magnetic_field_cgs,
             electric_field_cgs, electric_potential_cgs,
             resistance_cgs)

    def _convert(self, x, new_dims):
        old_dims = x.units.dimensions
        new_units, convert_factor = em_conversions[old_dims]
        return x.in_cgs().v*convert_factor, new_units

    def __str__(self):
        return "SI: EM CGS unit -> EM SI unit"

class ElectromagneticCGS(Equivalence):
    type_name = "CGS"
    alternate_names = ["cgs"]
    one_way = True
    _dims = (current_mks, charge_mks, magnetic_field_mks,
             electric_field_mks, electric_potential_mks,
             resistance_mks)

    def _convert(self, x, new_dims):
        old_dims = x.units.dimensions
        new_units, convert_factor = em_conversions[old_dims]
        return x.in_mks().v*convert_factor, new_units

    def __str__(self):
        return "CGS: EM SI unit -> EM CGS unit"
<|endoftext|>


<commit_before>__version__ = "0.2"
__license__ = "MIT"
<commit_msg>Version bump: 1.2<commit_after>__version__ = "1.2"
__license__ = "MIT"
<|endoftext|>