def CreateTask(self,
session_identifier,
storage_format=definitions.STORAGE_FORMAT_SQLITE):
"""Creates a task.
Args:
session_identifier (str): the identifier of the session the task is
part of.
storage_format (Optional[str]): the storage format that the task should be
stored in.
Returns:
Task: task attribute container.
"""
task = tasks.Task(session_identifier)
task.storage_format = storage_format
logger.debug('Created task: {0:s}.'.format(task.identifier))
with self._lock:
self._tasks_queued[task.identifier] = task
self._total_number_of_tasks += 1
self.SampleTaskStatus(task, 'created')
return task
def testScanForProcessedTasks(self):
"""Tests the ScanForProcessedTasks method"""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
# There should be no processed task identifiers initially.
task_identifiers, _ = redis_store.RedisStore.ScanForProcessedTasks(
session.identifier, redis_client=redis_client)
self.assertEqual([], task_identifiers)
# Opening and closing a writer for a task should cause the task to be marked
# as complete.
storage_writer = writer.RedisStorageWriter(
storage_type=definitions.STORAGE_TYPE_TASK)
storage_writer.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
storage_writer.Close()
# The task store is now marked as processed.
task_identifiers, _ = redis_store.RedisStore.ScanForProcessedTasks(
session.identifier, redis_client=redis_client)
self.assertEqual([task.identifier], task_identifiers)
def testCreateTaskStart(self):
"""Tests the CreateTaskStart function."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
task = tasks.Task(session_identifier=session_identifier)
task_start = task.CreateTaskStart()
self.assertIsNotNone(task_start)
def testHasAttributeContainers(self):
"""Tests the HasAttributeContainers method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertFalse(result)
test_store.AddAttributeContainer(event_data_stream)
result = test_store.HasAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertTrue(result)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetRedisHashName(self):
"""Tests the _GetRedisHashName function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
event_data_stream = events.EventDataStream()
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
redis_hash_name = test_store._GetRedisHashName(
event_data_stream.CONTAINER_TYPE)
expected_redis_hash_name = '{0:s}-{1:s}-{2:s}'.format(
task.session_identifier, task.identifier,
event_data_stream.CONTAINER_TYPE)
self.assertEqual(redis_hash_name, expected_redis_hash_name)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetSerializedAttributeContainers(self):
"""Tests the GetSerializedAttributeContainers method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
for _, event_data, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event_data)
cursor, serialized_containers = (
test_store.GetSerializedAttributeContainers(
'event_data', 0, 0))
self.assertEqual(len(serialized_containers), 4)
for serialized_container in serialized_containers:
self.assertIsInstance(serialized_container, bytes)
self.assertIsInstance(cursor, int)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetSortedEvents(self):
"""Tests the GetSortedEvents method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
for event, _, _ in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
test_store.AddAttributeContainer(event)
retrieved_events = list(test_store.GetSortedEvents())
self.assertEqual(len(retrieved_events), 4)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetNumberOfAttributeContainers(self):
"""Tests the GetNumberOfAttributeContainers function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
test_store.AddAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainerByIndex(self):
"""Tests the GetAttributeContainerByIndex function."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore(
storage_type=definitions.STORAGE_TYPE_TASK)
test_store.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNone(container)
test_store.AddAttributeContainer(event_data_stream)
container = test_store.GetAttributeContainerByIndex(
event_data_stream.CONTAINER_TYPE, 0)
self.assertIsNotNone(container)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testGetAttributeContainerByIdentifier(self):
"""Tests the GetAttributeContainerByIdentifier method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore()
test_store.Open(
redis_client=redis_client, session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
test_store.AddAttributeContainer(event_data_stream)
identifier = event_data_stream.GetIdentifier()
container = test_store.GetAttributeContainerByIdentifier(
event_data_stream.CONTAINER_TYPE, identifier)
self.assertIsNotNone(container)
identifier.sequence_number = 99
container = test_store.GetAttributeContainerByIdentifier(
event_data_stream.CONTAINER_TYPE, identifier)
self.assertIsNone(container)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
def testReadAndWriteSerializedTask(self):
"""Test ReadSerialized and WriteSerialized of Task."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
expected_task = tasks.Task(session_identifier=session_identifier)
json_string = (json_serializer.JSONAttributeContainerSerializer.
WriteSerialized(expected_task))
self.assertIsNotNone(json_string)
task = json_serializer.JSONAttributeContainerSerializer.ReadSerialized(
json_string)
self.assertIsNotNone(task)
self.assertIsInstance(task, tasks.Task)
expected_task_dict = {
'aborted': False,
'has_retry': False,
'identifier': task.identifier,
'session_identifier': session_identifier,
'start_time': task.start_time
}
task_dict = task.CopyToDict()
self.assertEqual(sorted(task_dict.items()),
sorted(expected_task_dict.items()))
def testCreateTaskCompletion(self):
"""Tests the CreateTaskCompletion function."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
task = tasks.Task(session_identifier=session_identifier)
task_completion = task.CreateTaskCompletion()
self.assertIsNotNone(task_completion)
def testMarkTaskAsMerging(self):
"""Tests the MarkTaskAsMerging method"""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
# Trying to mark a task as merging without finalizing it raises an error.
with self.assertRaises(IOError):
redis_store.RedisStore.MarkTaskAsMerging(task.identifier,
session.identifier,
redis_client=redis_client)
# Opening and closing a writer for a task should cause the task to be marked
# as complete.
storage_writer = writer.RedisStorageWriter(
storage_type=definitions.STORAGE_TYPE_TASK)
storage_writer.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
storage_writer.Close()
redis_store.RedisStore.MarkTaskAsMerging(task.identifier,
session.identifier,
redis_client=redis_client)
def testWriteTaskStartAndCompletion(self):
"""Tests the WriteTaskStart and WriteTaskCompletion functions."""
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
storage_writer = fake_writer.FakeStorageWriter(
storage_type=definitions.STORAGE_TYPE_TASK)
storage_writer.Open()
storage_writer.WriteTaskStart(task)
storage_writer.WriteTaskCompletion(task)
storage_writer.Close()
with self.assertRaises(IOError):
storage_writer.WriteTaskStart(task)
with self.assertRaises(IOError):
storage_writer.WriteTaskCompletion(task)
storage_writer = fake_writer.FakeStorageWriter()
storage_writer.Open()
with self.assertRaises(IOError):
storage_writer.WriteTaskStart(task)
with self.assertRaises(IOError):
storage_writer.WriteTaskCompletion(task)
storage_writer.Close()
def _CreateTaskStorageFile(self, session, path, event_values_list):
"""Creates a task storage file for testing.
Args:
session (Session): session the task storage is part of.
path (str): path to the task storage file that should be merged.
event_values_list (list[dict[str, str]]): list of event values.
"""
task = tasks.Task(session_identifier=session.identifier)
storage_file = writer.SQLiteStorageFileWriter(
session,
path,
storage_type=definitions.STORAGE_TYPE_TASK,
task=task)
storage_file.Open()
for event, event_data in containers_test_lib.CreateEventsFromValues(
event_values_list):
storage_file.AddEventData(event_data)
event.SetEventDataIdentifier(event_data.GetIdentifier())
storage_file.AddEvent(event)
storage_file.Close()
def _CreateTaskStore(self, session, redis_client):
"""Creates a task store for testing.
Args:
session (Session): session the task store is part of.
redis_client (Redis): Redis client to query. If specified, no
new client will be created.
Returns:
Task: the the task for the store that was created.
"""
task = tasks.Task(session_identifier=session.identifier)
task_store = redis_writer.RedisStorageWriter(
session, storage_type=definitions.STORAGE_TYPE_TASK, task=task)
task_store.Open(redis_client=redis_client)
for event, event_data in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
task_store.AddEventData(event_data)
event.SetEventDataIdentifier(event_data.GetIdentifier())
task_store.AddEvent(event)
task_store.Close()
return task
def testAddAttributeContainer(self):
"""Tests the AddAttributeContainer method."""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_store = redis_store.RedisStore()
test_store.Open(
redis_client=redis_client, session_identifier=task.session_identifier,
task_identifier=task.identifier)
try:
event_data_stream = events.EventDataStream()
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 0)
test_store.AddAttributeContainer(event_data_stream)
number_of_containers = test_store.GetNumberOfAttributeContainers(
event_data_stream.CONTAINER_TYPE)
self.assertEqual(number_of_containers, 1)
finally:
test_store.Close()
self._RemoveSessionData(redis_client, session.identifier)
with self.assertRaises(IOError):
test_store.AddAttributeContainer(event_data_stream)
def testContains(self):
"""Tests the __contains__ function."""
task = tasks.Task()
task.storage_file_size = 10
heap = task_manager._PendingMergeTaskHeap()
heap.PushTask(task)
self.assertIn(task.identifier, heap)
def testUpdateProcessingTime(self):
"""Tests the UpdateProcessingTime function."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
task = tasks.Task(session_identifier=session_identifier)
self.assertIsNone(task.last_processing_time)
task.UpdateProcessingTime()
self.assertIsNotNone(task.last_processing_time)
def testInitialization(self):
"""Tests the __init__ function."""
test_redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
test_reader = reader.RedisStorageReader(
session.identifier, task.identifier, redis_client=test_redis_client)
self.assertIsNotNone(test_reader)
def testLength(self):
"""Tests the __len__ function."""
task = tasks.Task()
task.storage_file_size = 10
heap = task_manager._PendingMergeTaskHeap()
self.assertEqual(len(heap), 0)
heap.PushTask(task)
self.assertEqual(len(heap), 1)
def testCreateRetryTask(self):
"""Tests the CreateRetryTask function."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
task = tasks.Task(session_identifier=session_identifier)
task.path_spec = 'test_path_spec'
retry_task = task.CreateRetryTask()
self.assertNotEqual(retry_task.identifier, task.identifier)
self.assertTrue(task.has_retry)
self.assertFalse(retry_task.has_retry)
self.assertEqual(retry_task.path_spec, task.path_spec)
def testPushTask(self):
"""Tests the PushTask function."""
task = tasks.Task()
task.storage_file_size = 10
heap = task_manager._PendingMergeTaskHeap()
self.assertEqual(len(heap), 0)
heap.PushTask(task)
self.assertEqual(len(heap), 1)
task = tasks.Task()
task.storage_file_size = 100
heap.PushTask(task)
self.assertEqual(len(heap), 2)
task = tasks.Task()
with self.assertRaises(ValueError):
heap.PushTask(task)
def testCreateRetry(self):
"""Tests the CreateRetry function."""
session_identifier = '{0:s}'.format(uuid.uuid4().hex)
task = tasks.Task(session_identifier=session_identifier)
task.path_spec = 'test_pathspec_value'
retry_task = task.CreateRetry()
self.assertEqual(task.path_spec, retry_task.path_spec)
self.assertNotEqual(task.identifier, retry_task.identifier)
self.assertTrue(task.retried)
self.assertEqual(task.identifier, retry_task.original_task_identifier)
def testTaskIsRetriable(self):
"""Tests the _TaskIsRetriable function."""
task = tasks.Task()
manager = task_manager.TaskManager()
result = manager._TaskIsRetriable(task)
self.assertTrue(result)
task.retried = True
result = manager._TaskIsRetriable(task)
self.assertFalse(result)
def testGetTaskPendingMerge(self):
"""Tests the GetTaskPendingMerge function."""
current_task = tasks.Task()
current_task.storage_file_size = 10
manager = task_manager.TaskManager()
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge when _tasks_pending_merge is
# empty.
result_task = manager.GetTaskPendingMerge(current_task)
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 0)
manager.UpdateTaskAsPendingMerge(task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge when there is a current task
# with a higher merge priority.
current_task.merge_priority = 1
task.merge_priority = 1000
result_task = manager.GetTaskPendingMerge(current_task)
self.assertIsNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 1)
self.assertEqual(len(manager._tasks_merging), 0)
# Determine if there are tasks pending merge.
result_task = manager.GetTaskPendingMerge(None)
self.assertIsNotNone(result_task)
self.assertEqual(len(manager._tasks_queued), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_merging), 1)
def CreateTask(self, session_identifier):
"""Creates a task.
Args:
session_identifier (str): the identifier of the session the task is
part of.
Returns:
Task: task attribute container.
"""
with self._lock:
task = tasks.Task(session_identifier)
self._active_tasks[task.identifier] = task
return task
def testGetProcessedTaskByIdentifier(self):
"""Tests the GetProcessedTaskByIdentifier function."""
manager = task_manager.TaskManager()
# Test with queued task.
test_tasks = []
for _ in range(3):
task = manager.CreateTask(self._TEST_SESSION_IDENTIFIER)
task.storage_file_size = 10
test_tasks.append(task)
self.assertEqual(len(manager._tasks_queued), 3)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
task = test_tasks[0]
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test with processing task.
manager.UpdateTaskAsProcessingByIdentifier(task.identifier)
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 1)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 0)
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test with abandoned task.
task.last_processing_time -= (2 * manager._TASK_INACTIVE_TIME *
definitions.MICROSECONDS_PER_SECOND)
manager._AbandonInactiveProcessingTasks()
self.assertEqual(len(manager._tasks_queued), 2)
self.assertEqual(len(manager._tasks_processing), 0)
self.assertEqual(len(manager._tasks_pending_merge), 0)
self.assertEqual(len(manager._tasks_abandoned), 1)
processed_task = manager.GetProcessedTaskByIdentifier(task.identifier)
self.assertEqual(processed_task, task)
# Test status of task is unknown.
task = tasks.Task()
with self.assertRaises(KeyError):
manager.GetProcessedTaskByIdentifier(task.identifier)
def testProcessTask(self):
"""Tests the _ProcessTask function."""
session = sessions.Session()
storage_writer = self._CreateStorageWriter(session)
knowledge_base = self._CreateKnowledgeBase()
configuration = configurations.ProcessingConfiguration()
test_process = worker_process.WorkerProcess(
None, storage_writer, knowledge_base, session.identifier, configuration,
name='TestWorker')
test_process._parser_mediator = self._CreateParserMediator(
storage_writer, knowledge_base)
task = tasks.Task(session_identifier=session.identifier)
test_process._ProcessTask(task)
def testPopTask(self):
"""Tests the PopTask function."""
task = tasks.Task()
task.storage_file_size = 10
heap = task_manager._PendingMergeTaskHeap()
result_task = heap.PopTask()
self.assertIsNone(result_task)
heap.PushTask(task)
self.assertEqual(len(heap), 1)
result_task = heap.PopTask()
self.assertEqual(len(heap), 0)
self.assertEqual(result_task, task)
