6 files changed, 317 insertions, 0 deletions
diff --git a/benchmarks/harness/machine_producers/BUILD b/benchmarks/harness/machine_producers/BUILD
new file mode 100644
index 000000000..a48da02a1
--- /dev/null
+++ b/benchmarks/harness/machine_producers/BUILD
@@ -0,0 +1,40 @@
+load("//benchmarks:defs.bzl", "py_library", "requirement")
+
+package(
+    default_visibility = ["//benchmarks:__subpackages__"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "harness",
+    srcs = ["__init__.py"],
+)
+
+py_library(
+    name = "machine_producer",
+    srcs = ["machine_producer.py"],
+)
+
+py_library(
+    name = "mock_producer",
+    srcs = ["mock_producer.py"],
+    deps = [
+        "//benchmarks/harness:machine",
+        "//benchmarks/harness/machine_producers:machine_producer",
+    ],
+)
+
+py_library(
+    name = "yaml_producer",
+    srcs = ["yaml_producer.py"],
+    deps = [
+        "//benchmarks/harness:machine",
+        "//benchmarks/harness/machine_producers:machine_producer",
+        requirement("PyYAML", False),
+    ],
+)
+
+py_library(
+    name = "gcloud_mock_recorder",
+    srcs = ["gcloud_mock_recorder.py"],
+)
diff --git a/benchmarks/harness/machine_producers/__init__.py b/benchmarks/harness/machine_producers/__init__.py
new file mode 100644
index 000000000..634ef4843
--- /dev/null
+++ b/benchmarks/harness/machine_producers/__init__.py
@@ -0,0 +1,13 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/benchmarks/harness/machine_producers/gcloud_mock_recorder.py b/benchmarks/harness/machine_producers/gcloud_mock_recorder.py
new file mode 100644
index 000000000..fd9837a37
--- /dev/null
+++ b/benchmarks/harness/machine_producers/gcloud_mock_recorder.py
@@ -0,0 +1,97 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A recorder and replay for testing the GCloudProducer.
+
+MockPrinter and MockReader handle printing and reading mock data for the
+purposes of testing. MockPrinter is passed to GCloudProducer objects. The user
+can then run scenarios and record them for playback in tests later.
+
+MockReader is passed to MockGcloudProducer objects and handles reading the
+previously recorded mock data.
+
+It is left to the user to check if data printed is properly redacted for their
+own use. The intended usecase for this class is data coming from gcloud
+commands, which will contain public IPs and other instance data.
+
+The data format is json and printed/read from the ./test_data directory. The
+data is the output of subprocess.CompletedProcess objects in json format.
+
+  Typical usage example:
+
+  recorder = MockPrinter()
+  producer = GCloudProducer(args, recorder)
+  machines = producer.get_machines(1)
+  with open("my_file.json") as fd:
+    recorder.write_out(fd)
+
+  reader = MockReader(filename)
+  producer = MockGcloudProducer(args, mock)
+  machines = producer.get_machines(1)
+  assert len(machines) == 1
+"""
+
+import io
+import json
+import subprocess
+
+
+class MockPrinter(object):
+  """Handles printing Mock data for MockGcloudProducer.
+
+  Attributes:
+    _records: list of json object records for printing
+  """
+
+  def __init__(self):
+    self._records = []
+
+  def record(self, entry: subprocess.CompletedProcess):
+    """Records data and strips out ip addresses."""
+
+    record = {
+        "args": entry.args,
+        "stdout": entry.stdout.decode("utf-8"),
+        "returncode": str(entry.returncode)
+    }
+    self._records.append(record)
+
+  def write_out(self, fd: io.FileIO):
+    """Prints out the data into the given filepath."""
+    fd.write(json.dumps(self._records, indent=4))
+
+
+class MockReader(object):
+  """Handles reading Mock data for MockGcloudProducer.
+
+  Attributes:
+    _records: List[json] records read from the passed in file.
+  """
+
+  def __init__(self, filepath: str):
+    with open(filepath, "rb") as file:
+      self._records = json.loads(file.read())
+      self._i = 0
+
+  def __iter__(self):
+    return self
+
+  def __next__(self, args) -> subprocess.CompletedProcess:
+    """Returns the next record as a CompletedProcess."""
+    if self._i < len(self._records):
+      record = self._records[self._i]
+      stdout = record["stdout"].encode("ascii")
+      returncode = int(record["returncode"])
+      return subprocess.CompletedProcess(
+          args=args, returncode=returncode, stdout=stdout, stderr=b"")
+    raise StopIteration()
diff --git a/benchmarks/harness/machine_producers/machine_producer.py b/benchmarks/harness/machine_producers/machine_producer.py
new file mode 100644
index 000000000..124ee14cc
--- /dev/null
+++ b/benchmarks/harness/machine_producers/machine_producer.py
@@ -0,0 +1,30 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Abstract types."""
+
+from typing import List
+
+from benchmarks.harness import machine
+
+
+class MachineProducer:
+  """Abstract Machine producer."""
+
+  def get_machines(self, num_machines: int) -> List[machine.Machine]:
+    """Returns the requested number of machines."""
+    raise NotImplementedError
+
+  def release_machines(self, machine_list: List[machine.Machine]):
+    """Releases the given set of machines."""
+    raise NotImplementedError
diff --git a/benchmarks/harness/machine_producers/mock_producer.py b/benchmarks/harness/machine_producers/mock_producer.py
new file mode 100644
index 000000000..4f29ad53f
--- /dev/null
+++ b/benchmarks/harness/machine_producers/mock_producer.py
@@ -0,0 +1,31 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Producers of mocks."""
+
+from typing import List
+
+from benchmarks.harness import machine
+from benchmarks.harness.machine_producers import machine_producer
+
+
+class MockMachineProducer(machine_producer.MachineProducer):
+  """Produces MockMachine objects."""
+
+  def get_machines(self, num_machines: int) -> List[machine.MockMachine]:
+    """Returns the request number of MockMachines."""
+    return [machine.MockMachine() for i in range(num_machines)]
+
+  def release_machines(self, machine_list: List[machine.MockMachine]):
+    """No-op."""
+    return
diff --git a/benchmarks/harness/machine_producers/yaml_producer.py b/benchmarks/harness/machine_producers/yaml_producer.py
new file mode 100644
index 000000000..5d334e480
--- /dev/null
+++ b/benchmarks/harness/machine_producers/yaml_producer.py
@@ -0,0 +1,106 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Producers based on yaml files."""
+
+import os
+import threading
+from typing import Dict
+from typing import List
+
+import yaml
+
+from benchmarks.harness import machine
+from benchmarks.harness.machine_producers import machine_producer
+
+
+class YamlMachineProducer(machine_producer.MachineProducer):
+  """Loads machines from a yaml file."""
+
+  def __init__(self, path: str):
+    self.machines = build_machines(path)
+    self.max_machines = len(self.machines)
+    self.machine_condition = threading.Condition()
+
+  def get_machines(self, num_machines: int) -> List[machine.Machine]:
+    if num_machines > self.max_machines:
+      raise ValueError(
+          "Insufficient Ammount of Machines. {ask} asked for and have {max_num} max."
+          .format(ask=num_machines, max_num=self.max_machines))
+
+    with self.machine_condition:
+      while not self._enough_machines(num_machines):
+        self.machine_condition.wait(timeout=1)
+      return [self.machines.pop(0) for _ in range(num_machines)]
+
+  def release_machines(self, machine_list: List[machine.Machine]):
+    with self.machine_condition:
+      while machine_list:
+        next_machine = machine_list.pop()
+        self.machines.append(next_machine)
+      self.machine_condition.notify()
+
+  def _enough_machines(self, ask: int):
+    return ask <= len(self.machines)
+
+
+def build_machines(path: str, num_machines: str = -1) -> List[machine.Machine]:
+  """Builds machine objects defined by the yaml file "path".
+
+  Args:
+    path: The path to a yaml file which defines machines.
+    num_machines: Optional limit on how many machine objects to build.
+
+  Returns:
+    Machine objects in a list.
+
+    If num_machines is set, len(machines) <= num_machines.
+  """
+  data = parse_yaml(path)
+  machines = []
+  for key, value in data.items():
+    if len(machines) == num_machines:
+      return machines
+    if isinstance(value, dict):
+      machines.append(machine.RemoteMachine(key, **value))
+    else:
+      machines.append(machine.LocalMachine(key))
+  return machines
+
+
+def parse_yaml(path: str) -> Dict[str, Dict[str, str]]:
+  """Parse the yaml file pointed by path.
+
+  Args:
+    path: The path to yaml file.
+
+  Returns:
+    The contents of the yaml file as a dictionary.
+  """
+  data = get_file_contents(path)
+  return yaml.load(data, Loader=yaml.Loader)
+
+
+def get_file_contents(path: str) -> str:
+  """Dumps the file contents to a string and returns them.
+
+  Args:
+    path: The path to dump.
+
+  Returns:
+    The file contents as a string.
+  """
+  if not os.path.isabs(path):
+    path = os.path.abspath(path)
+  with open(path) as input_file:
+    return input_file.read()