<commit_before>import time
from collections import namedtuple
import threading
import mock

from django.utils import unittest

from chroma_agent.lib.shell import AgentShell
from chroma_agent.device_plugins.action_runner import ActionRunnerPlugin, CallbackAfterResponse
from chroma_agent.plugin_manager import ActionPluginManager, DevicePlugin
from chroma_agent.agent_client import AgentDaemonContext

ACTION_ONE_NO_CONTEXT_RETVAL = 'action_one_no_context_return'
ACTION_ONE_WITH_CONTEXT_RETVAL = 'action_one_with_context_return'
ACTION_TWO_RETVAL = 'action_two_return'


subprocesses = {
    ('subprocess_one', 'subprocess_one_arg'): lambda: AgentShell.RunResult(0, 'subprocess_one_stdout', 'subprocess_one_stderr', False),
    ('subprocess_two', 'subprocess_two_arg'): lambda: AgentShell.RunResult(-1, 'subprocess_two_stdout', 'subprocess_two_stderr', False)
}


def action_one_no_context(arg1):
    """An action which invokes subprocess_one"""

    assert arg1 == "arg1_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    return ACTION_ONE_NO_CONTEXT_RETVAL


def action_one_with_context(agent_daemon_context, arg1):
    """An action which invokes subprocess_one"""

    assert isinstance(agent_daemon_context, AgentDaemonContext)
    assert arg1 == "arg1_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    return ACTION_ONE_WITH_CONTEXT_RETVAL


def action_two(arg1):
    """An action which invokes subprocess_one and subprocess_two"""

    assert arg1 == "arg2_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    AgentShell.try_run(['subprocess_two', 'subprocess_two_arg'])
    return ACTION_TWO_RETVAL


def action_three():
    """An action which runs for a long time unless interrupted"""

    AgentShell.try_run(['sleep', '10000'])
    raise AssertionError("subprocess three should last forever!")


def action_four():
    """An action which raises an CallbackAfterResponse"""

    def action_four_callback():
        return 'action_four_called_back'

    raise CallbackAfterResponse('result', action_four_callback)

ACTIONS = [action_one_no_context, action_one_with_context, action_two, action_three, action_four]


class ActionTestCase(unittest.TestCase):
    MOCK_SUBPROCESSES = None

    SendMessageParams = namedtuple('SendMessageParams', ['body', 'callback'])

    def mock_send_message(self, body, callback = None):
        self.mock_send_message_lock.acquire()
        self.mock_send_message_call_args_list.append(self.SendMessageParams(body, callback))
        self.mock_send_message_call_count += 1
        self.mock_send_message_lock.release()

    def mock_run(self, args, logger, result_store, timeout):
        return self.MOCK_SUBPROCESSES[tuple(args)]()

    def setUp(self):
        # Register some testing actions
        self.action_plugins = ActionPluginManager()
        for action in ACTIONS:
            self.action_plugins.commands[action.__name__] = action

        # Intercept messages sent with a thread safe mock<commit_msg>Add extra argument to mock (#248)<commit_after>import time
from collections import namedtuple
import threading
import mock

from django.utils import unittest

from chroma_agent.lib.shell import AgentShell
from chroma_agent.device_plugins.action_runner import ActionRunnerPlugin, CallbackAfterResponse
from chroma_agent.plugin_manager import ActionPluginManager, DevicePlugin
from chroma_agent.agent_client import AgentDaemonContext

ACTION_ONE_NO_CONTEXT_RETVAL = 'action_one_no_context_return'
ACTION_ONE_WITH_CONTEXT_RETVAL = 'action_one_with_context_return'
ACTION_TWO_RETVAL = 'action_two_return'


subprocesses = {
    ('subprocess_one', 'subprocess_one_arg'): lambda: AgentShell.RunResult(0, 'subprocess_one_stdout', 'subprocess_one_stderr', False),
    ('subprocess_two', 'subprocess_two_arg'): lambda: AgentShell.RunResult(-1, 'subprocess_two_stdout', 'subprocess_two_stderr', False)
}


def action_one_no_context(arg1):
    """An action which invokes subprocess_one"""

    assert arg1 == "arg1_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    return ACTION_ONE_NO_CONTEXT_RETVAL


def action_one_with_context(agent_daemon_context, arg1):
    """An action which invokes subprocess_one"""

    assert isinstance(agent_daemon_context, AgentDaemonContext)
    assert arg1 == "arg1_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    return ACTION_ONE_WITH_CONTEXT_RETVAL


def action_two(arg1):
    """An action which invokes subprocess_one and subprocess_two"""

    assert arg1 == "arg2_test"
    stdout = AgentShell.try_run(['subprocess_one', 'subprocess_one_arg'])
    assert stdout == 'subprocess_one_stdout'
    AgentShell.try_run(['subprocess_two', 'subprocess_two_arg'])
    return ACTION_TWO_RETVAL


def action_three():
    """An action which runs for a long time unless interrupted"""

    AgentShell.try_run(['sleep', '10000'])
    raise AssertionError("subprocess three should last forever!")


def action_four():
    """An action which raises an CallbackAfterResponse"""

    def action_four_callback():
        return 'action_four_called_back'

    raise CallbackAfterResponse('result', action_four_callback)

ACTIONS = [action_one_no_context, action_one_with_context, action_two, action_three, action_four]


class ActionTestCase(unittest.TestCase):
    MOCK_SUBPROCESSES = None

    SendMessageParams = namedtuple('SendMessageParams', ['body', 'callback'])

    def mock_send_message(self, body, callback = None):
        self.mock_send_message_lock.acquire()
        self.mock_send_message_call_args_list.append(self.SendMessageParams(body, callback))
        self.mock_send_message_call_count += 1
        self.mock_send_message_lock.release()

    def mock_run(self, args, logger, result_store, timeout, shell=False):
        return self.MOCK_SUBPROCESSES[tuple(args)]()

    def setUp(self):
        # Register some testing actions
        self.action_plugins = ActionPluginManager()
        for action in ACTIONS:
            self.action_plugins.commands[action.__name__] = action

        # Intercept messages sent with a thread safe mock
{'fraction_non_alphanumeric': 0.08963367108339829, 'fraction_numerical': 0.005767731878409977, 'mean_word_length': 5.396809571286141, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '2977229', 'n_tokens_mistral': 1965, 'n_tokens_neox': 1912, 'n_words': 415}