Repository: google/openhtf

Issue:
Provide an option to disable source code capturing.
Some test authors might consider their tester code "IP" and want to prevent third party CM's from easily viewing it. We should consider having an option to exclude source code from the test record so it's not transmitted in plaintext over the test network, and it's not displayed in the operator interface (frontend).

For the record, I personally don't believe this is a good way to address the imagined problem, but I've had a few requests from test authors provide this option.



Discussion:

Resolving patch:
diff --git a/openhtf/__init__.py b/openhtf/__init__.py
--- a/openhtf/__init__.py
+++ b/openhtf/__init__.py
@@ -26,6 +26,7 @@
 import signal
 import socket
 import sys
+import textwrap
 import threading
 import weakref
 from types import LambdaType
@@ -54,6 +55,14 @@
 __version__ = util.get_version()
 _LOG = logging.getLogger(__name__)
 
+conf.declare('capture_source', description=textwrap.dedent(
+    '''Whether to capture the source of phases and the test module.  This
+    defaults to False since this potentially reads many files and makes large
+    string copies.
+
+    Set to 'true' if you want to capture your test's source.'''),
+    default_value=False)
+
 
 class UnrecognizedTestUidError(Exception):
   """Raised when information is requested about an unknown Test UID."""
@@ -94,11 +103,24 @@ def __init__(self, *phases, **metadata):
 
     self.created_time_millis = util.time_millis()
     self.last_run_time_millis = None
-    code_info = test_record.CodeInfo.for_module_from_stack(levels_up=2)
-    self._test_desc = TestDescriptor(self.uid, phases, code_info, metadata)
     self._test_options = TestOptions()
     self._lock = threading.Lock()
     self._executor = None
+    self._test_desc = TestDescriptor(
+        self.uid, phases, test_record.CodeInfo.uncaptured(), metadata)
+
+    if conf.capture_source:
+      # First, we copy the phases with the real CodeInfo for them.
+      phases = [
+        mutablerecords.CopyRecord(
+            phase, code_info=test_record.CodeInfo.for_function(phase.func))
+        for phase in self._test_desc.phases]
+
+      # Then we replace the TestDescriptor with one that stores the test
+      # module's CodeInfo as well as our newly copied phases.
+      code_info = test_record.CodeInfo.for_module_from_stack(levels_up=2)
+      self._test_desc = self._test_desc._replace(
+          code_info=code_info, phases=phases)
 
     # Make sure configure() gets called at least once before Execute().  The
     # user might call configure() again to override options, but we don't want
@@ -159,14 +181,7 @@ def add_output_callbacks(self, *callbacks):
     self._test_options.output_callbacks.extend(callbacks)
 
   def configure(self, **kwargs):
-    """Update test-wide configuration options.
-
-    Valid kwargs:
-      output_callbacks: List of output callbacks to run, typically it's better
-          to use add_output_callbacks(), but you can pass [] here to reset them.
-      teardown_function: Function to run at teardown.  We pass the same
-          arguments to it as a phase.
-    """
+    """Update test-wide configuration options. See TestOptions for docs."""
     # These internally ensure they are safe to call multiple times with no weird
     # side effects.
     create_arg_parser(add_help=True).parse_known_args()
@@ -227,6 +242,9 @@ def trigger_phase(test):
       else:
         trigger = test_start
 
+      if conf.capture_source:
+        trigger.code_info = test_record.CodeInfo.for_function(trigger.func)
+
       self._executor = core.TestExecutor(
           self._test_desc, trigger, self._test_options.teardown_function)
       _LOG.info('Executing test: %s', self.descriptor.code_info.name)
@@ -259,7 +277,14 @@ class TestOptions(mutablerecords.Record('TestOptions', [], {
     'output_callbacks': list,
     'teardown_function': None,
 })):
-  """Class encapsulating various tunable knobs for Tests and their defaults."""
+  """Class encapsulating various tunable knobs for Tests and their defaults.
+
+  name: The name of the test to be put into the metadata.
+  output_callbacks: List of output callbacks to run, typically it's better to
+      use add_output_callbacks(), but you can pass [] here to reset them.
+  teardown_function: Function to run at teardown.  We pass the same arguments to
+      it as a phase.
+  """
 
 
 class TestDescriptor(collections.namedtuple(
@@ -366,17 +391,18 @@ class PhasePlug(mutablerecords.Record(
 
 
 class PhaseDescriptor(mutablerecords.Record(
-    'PhaseDescriptor', ['func', 'code_info'],
+    'PhaseDescriptor', ['func'],
     {'options': PhaseOptions, 'plugs': list, 'measurements': list,
-     'extra_kwargs': dict})):
+     'extra_kwargs': dict, 'code_info': test_record.CodeInfo.uncaptured()})):
   """Phase function and related information.
 
   Attributes:
     func: Function to be called (with TestApi as first argument).
-    code_info: Info about the source code of func.
     options: PhaseOptions instance.
     plugs: List of PhasePlug instances.
     measurements: List of Measurement objects.
+    extra_kwargs: Keyword arguments that will be passed to the function.
+    code_info: Info about the source code of func.
   """
 
   @classmethod
@@ -398,7 +424,7 @@ def wrap_or_copy(cls, func, **options):
       # or kwargs.  See with_args() below for more details.
       retval = mutablerecords.CopyRecord(func)
     else:
-      retval = cls(func, test_record.CodeInfo.for_function(func))
+      retval = cls(func)
     retval.options.update(**options)
     return retval
 


Test patch:
diff --git a/openhtf/core/test_record.py b/openhtf/core/test_record.py
--- a/openhtf/core/test_record.py
+++ b/openhtf/core/test_record.py
@@ -120,3 +120,7 @@ def for_module_from_stack(cls, levels_up=1):
   def for_function(cls, func):
     source = _get_source_safely(func)
     return cls(func.__name__, inspect.getdoc(func), source)
+
+  @classmethod
+  def uncaptured(cls):
+    return cls('', None, '')
diff --git a/test/capture_source_test.py b/test/capture_source_test.py
new file mode 100644
--- /dev/null
+++ b/test/capture_source_test.py
@@ -0,0 +1,41 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+
+import openhtf as htf
+
+
+def phase(test):
+  pass
+
+
+class BasicCodeCaptureTest(unittest.TestCase):
+  """Basic sanity test only for capture_source's behavior."""
+
+  @htf.conf.save_and_restore
+  def testCaptured(self):
+    htf.conf.load(capture_source=True)
+    test = htf.Test(phase)
+    phase_descriptor = test.descriptor.phases[0]
+    self.assertEquals(phase_descriptor.code_info.name, phase.__name__)
+
+  @htf.conf.save_and_restore
+  def testNotCaptured(self):
+    htf.conf.load(capture_source=False)
+    test = htf.Test(phase)
+    phase_descriptor = test.descriptor.phases[0]
+    self.assertEquals(phase_descriptor.code_info.name, '')
+
diff --git a/test/core/measurements_record.pickle b/test/core/measurements_record.pickle
Binary files a/test/core/measurements_record.pickle and b/test/core/measurements_record.pickle differ
diff --git a/test/core/measurements_test.py b/test/core/measurements_test.py
--- a/test/core/measurements_test.py
+++ b/test/core/measurements_test.py
@@ -96,7 +96,9 @@ def test_unit_enforcement(self):
     """Creating a measurement with invalid units should raise."""
     self.assertRaises(TypeError, htf.Measurement('bad_units').with_units, 1701)
 
-  @conf.save_and_restore(station_id='measurements_test', station_api_port=None)
+  @conf.save_and_restore(
+      station_id='measurements_test', station_api_port=None,
+      capture_source=True)
   def test_measurements(self):
     result = util.NonLocalResult()
     def _save_result(test_record):
diff --git a/test/util/conf_test.py b/test/util/conf_test.py
--- a/test/util/conf_test.py
+++ b/test/util/conf_test.py
@@ -96,6 +96,7 @@ def test_as_dict(self):
         'enable_station_discovery': True,
         'station_api_port': 8888,
         'allow_unset_measurements': False,
+        'capture_source': False,
         'station_discovery_string': 'OPENHTF_DISCOVERY',
         'station_api_bind_address': '0.0.0.0',
         'station_id': 'station_id',