<commit_before>#
# 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 os
import threading

import mock
from mesos.interface import mesos_pb2
from twitter.common.contextutil import temporary_dir
from twitter.common.quantity import Amount, Time

from apache.aurora.executor.common.resource_manager import ResourceManagerProvider
from apache.aurora.executor.common.sandbox import DirectorySandbox
from apache.thermos.monitoring.disk import DiskCollector


# TODO(jcohen): There should really be a single canonical source for creating test jobs/tasks
def make_assigned_task(thermos_config, assigned_ports=None):
  from gen.apache.aurora.api.constants import AURORA_EXECUTOR_NAME
  from gen.apache.aurora.api.ttypes import AssignedTask, ExecutorConfig, Identity, TaskConfig

  assigned_ports = assigned_ports or {}
  executor_config = ExecutorConfig(name=AURORA_EXECUTOR_NAME, data=thermos_config.json_dumps())
  task_config = TaskConfig(
    owner=Identity(role=thermos_config.role().get(), user=thermos_config.role().get()),
    environment=thermos_config.environment().get(),
    jobName=thermos_config.name().get(),
    executorConfig=executor_config)

  return AssignedTask(instanceId=12345, task=task_config, assignedPorts=assigned_ports)


def make_job(role, environment, name, primary_port, portmap):
  from apache.aurora.config.schema.base import (
    Announcer,
    Job,
    Process,
    Resources,
    Task,
  )
  task = Task(
    name='ignore2',
    processes=[Process(name='ignore3', cmdline='ignore4')],
    resources=Resources(cpu=1, ram=1, disk=1))
  job = Job(
    role=role,
    environment=environment,
    name=name,
    cluster='ignore1',
    task=task,
    announce=Announcer(primary_port=primary_port, portmap=portmap))
  return job


def test_resource_manager():
  with temporary_dir() as td:
    sandbox = os.path.join(td, 'sandbox')
    root = os.path.join(td, 'thermos')

    completed_event = threading.Event()
    completed_event.set()
    completed_mock = mock.PropertyMock(completed_event)
    mock_disk_collector = mock.create_autospec(DiskCollector, spec_set=True)
    mock_disk_collector.sample.return_value = None
    value_mock = mock.PropertyMock(return_value=4197)
    type(mock_disk_collector).value = value_mock
    type(mock_disk_collector).completed_event = completed_mock

    rmp = ResourceManagerProvider(
        root,
        disk_collector=lambda sandbox: mock_disk_collector,
        disk_collection_interval=Amount(1, Time.SECONDS))
    rm = rmp.from_assigned_task(
        make_assigned_task(
            make_job('some-role', 'some-env', 'some-job', 'http', portmap={'http': 80})),
        DirectorySandbox(sandbox))

    assert rm.status is None

    try:
      rm.start()
      while rm.status is None:
        rm._kill_event.wait(timeout=1)
      result = rm.status
      assert result is not None
      assert result.reason.startswith('Disk limit exceeded')<commit_msg>Minor clean up to mock usage in resource manager integration tests.<commit_after>#
# 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 os
import threading

import mock
from mesos.interface import mesos_pb2
from twitter.common.contextutil import temporary_dir
from twitter.common.quantity import Amount, Time

from apache.aurora.executor.common.resource_manager import ResourceManagerProvider
from apache.aurora.executor.common.sandbox import DirectorySandbox
from apache.thermos.monitoring.disk import DiskCollector


# TODO(jcohen): There should really be a single canonical source for creating test jobs/tasks
def make_assigned_task(thermos_config, assigned_ports=None):
  from gen.apache.aurora.api.constants import AURORA_EXECUTOR_NAME
  from gen.apache.aurora.api.ttypes import AssignedTask, ExecutorConfig, Identity, TaskConfig

  assigned_ports = assigned_ports or {}
  executor_config = ExecutorConfig(name=AURORA_EXECUTOR_NAME, data=thermos_config.json_dumps())
  task_config = TaskConfig(
    owner=Identity(role=thermos_config.role().get(), user=thermos_config.role().get()),
    environment=thermos_config.environment().get(),
    jobName=thermos_config.name().get(),
    executorConfig=executor_config)

  return AssignedTask(instanceId=12345, task=task_config, assignedPorts=assigned_ports)


def make_job(role, environment, name, primary_port, portmap):
  from apache.aurora.config.schema.base import (
    Announcer,
    Job,
    Process,
    Resources,
    Task,
  )
  task = Task(
    name='ignore2',
    processes=[Process(name='ignore3', cmdline='ignore4')],
    resources=Resources(cpu=1, ram=1, disk=1))
  job = Job(
    role=role,
    environment=environment,
    name=name,
    cluster='ignore1',
    task=task,
    announce=Announcer(primary_port=primary_port, portmap=portmap))
  return job


def test_resource_manager():
  with temporary_dir() as td:
    sandbox = os.path.join(td, 'sandbox')
    root = os.path.join(td, 'thermos')

    mock_disk_collector_class = mock.create_autospec(DiskCollector, spec_set=True)
    mock_disk_collector = mock_disk_collector_class.return_value

    mock_disk_collector.sample.return_value = None
    value_mock = mock.PropertyMock(return_value=4197)
    type(mock_disk_collector).value = value_mock

    completed_event = threading.Event()
    completed_event.set()
    completed_mock = mock.PropertyMock(completed_event)
    type(mock_disk_collector).completed_event = completed_mock

    rmp = ResourceManagerProvider(
        root,
        disk_collector=mock_disk_collector_class,
        disk_collection_interval=Amount(1, Time.SECONDS))
    rm = rmp.from_assigned_task(
        make_assigned_task(
            make_job('some-role', 'some-env', 'some-job', 'http', portmap={'http': 80})),
        DirectorySandbox(sandbox))

    assert rm.status is None

    try:
      rm.start()
      while rm.status is None:
        rm._kill_event.wait(timeout=1)
      result = rm.status
      assert result is not None
      assert result.reason.startswith('Disk limit exceeded')
{'fraction_non_alphanumeric': 0.09280275822439304, 'fraction_numerical': 0.007182876023559833, 'mean_word_length': 4.490536277602524, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '379848', 'n_tokens_mistral': 2145, 'n_tokens_neox': 2080, 'n_words': 519}