Improve Worker Package Test Coverage and Fix Pre-existing Tests

tests/test_worker/test_pulse_tasks.py

@@ -1,16 +1,17 @@
+import copy
 from threading import local
 
 import pytest
+from celery.exceptions import Retry
 
 from treeherder.etl.exceptions import MissingPushError
 from treeherder.etl.push import store_push_data
 from treeherder.etl.tasks.pulse_tasks import store_pulse_tasks
 from treeherder.model.models import Job
 
 
-@pytest.mark.skip("Test needs fixing in bug: 1307289 (plus upgrade from jobs to tasks)")
 def test_retry_missing_revision_succeeds(
-    sample_data, sample_push, test_repository, mock_log_parser, monkeypatch
+    sample_data, sample_push, test_repository, mock_log_parser, failure_classifications, monkeypatch
 ):
     """
     Ensure that when the missing push exists after a retry, that the job
@@ -19,19 +20,37 @@ def test_retry_missing_revision_succeeds(
     thread_data = local()
     thread_data.retries = 0
     rs = sample_push[0]
-    job = sample_data.pulse_jobs[0]
+    job = copy.deepcopy(sample_data.pulse_jobs[0])
     job["origin"]["revision"] = rs["revision"]
     job["origin"]["project"] = test_repository.name
 
+    # Mock handle_message inside store_pulse_tasks so we don't hit the network, and can raise MissingPushError on first call
+    async def mock_handle_message(message, task_definition=None):
+        if thread_data.retries == 0:
+            raise MissingPushError("Missing push!")
+
+        return [job]
+
+    monkeypatch.setattr("treeherder.etl.tasks.pulse_tasks.handle_message", mock_handle_message)
+
     orig_retry = store_pulse_tasks.retry
 
-    def retry_mock(exc=None, countdown=None):
+    def retry_mock(exc=None, countdown=None, *args, **kwargs):
         assert isinstance(exc, MissingPushError)
         thread_data.retries += 1
         store_push_data(test_repository, [rs])
-        return orig_retry(exc=exc, countdown=countdown)
+        return orig_retry(exc=exc, countdown=countdown, *args, **kwargs)
 
     monkeypatch.setattr(store_pulse_tasks, "retry", retry_mock)
+
+    # First attempt should raise Retry because push is missing
+    with pytest.raises(Retry):
+        store_pulse_tasks.delay(job, "foo", "bar")
+
+    assert thread_data.retries == 1
+    assert Job.objects.count() == 0
+
+    # Second attempt should succeed because push is now stored
     store_pulse_tasks.delay(job, "foo", "bar")
 
     assert Job.objects.count() == 1

tests/test_worker/test_stats.py

@@ -5,7 +5,7 @@
 from django.utils import timezone
 
 from treeherder.model.models import Job, Push
-from treeherder.workers.stats import publish_stats
+from treeherder.workers.stats import get_stats_client, publish_stats
 
 
 @pytest.mark.django_db
@@ -17,7 +17,11 @@ def test_publish_stats_nothing_to_do(get_worker_mock, django_assert_num_queries,
     assert Job.objects.count() == 0
     with django_assert_num_queries(2):
         publish_stats()
-    assert [(level, message) for _, level, message in caplog.record_tuples] == [
+    assert [
+        (level, message)
+        for name, level, message in caplog.record_tuples
+        if name == "treeherder.workers.stats"
+    ] == [
         (20, "Publishing runtime statistics to statsd"),
         (20, "Ingested 0 pushes"),
         (20, "Ingested 0 jobs in total"),
@@ -41,7 +45,11 @@ def test_publish_stats(
 
     with django_assert_num_queries(2):
         publish_stats()
-    assert [(level, message) for _, level, message in caplog.record_tuples] == [
+    assert [
+        (level, message)
+        for name, level, message in caplog.record_tuples
+        if name == "treeherder.workers.stats"
+    ] == [
         (20, "Publishing runtime statistics to statsd"),
         (20, "Ingested 22 pushes"),
         (20, "Ingested 11 jobs in total"),
@@ -52,3 +60,13 @@ def test_publish_stats(
         call("jobs_repo.mozilla-central", 11),
         call("jobs_state.completed", 11),
     ]
+
+
+def test_get_stats_client(settings):
+    "Test get_stats_client returns a statsd client with correct settings"
+    settings.STATSD_HOST = "localhost"
+    settings.STATSD_PORT = 8125
+    settings.STATSD_PREFIX = "test-prefix"
+    client = get_stats_client()
+    assert client._addr[1] == 8125
+    assert client._prefix == "test-prefix"

tests/test_worker/test_task.py

@@ -1,18 +1,20 @@
+import zlib
 from functools import wraps
 from threading import local
+from unittest.mock import patch
 
+import jsonschema
 import pytest
 from celery.exceptions import Retry
-from django.db.utils import OperationalError
+from django.db.utils import IntegrityError, OperationalError, ProgrammingError
 
+from treeherder.etl.exceptions import MissingPushError
 from treeherder.workers.task import retryable_task
 
 thread_data = local()
 
 
 def count_retries(f):
-    thread_data.retry_count = -1
-
     @wraps(f)
     def inner():
         thread_data.retry_count += 1
@@ -21,6 +23,12 @@ def inner():
     return inner
 
 
+@pytest.fixture(autouse=True)
+def reset_retry_count():
+    thread_data.retry_count = -1
+    yield
+
+
 @retryable_task()
 def successful_task(x, y):
     return x + y
@@ -35,15 +43,76 @@ def test_retryable_task():
 
 @retryable_task()
 @count_retries
-def throwing_task():
+def throwing_task_type_error():
     raise TypeError
 
 
-def test_retryable_task_throws():
-    "Test celery immediately raises an error for a task that throws"
+@retryable_task()
+@count_retries
+def throwing_task_key_error():
+    raise KeyError
+
+
+@retryable_task()
+@count_retries
+def throwing_task_value_error():
+    raise ValueError
+
+
+@retryable_task()
+@count_retries
+def throwing_task_index_error():
+    raise IndexError
+
+
+@retryable_task()
+@count_retries
+def throwing_task_integrity_error():
+    raise IntegrityError
+
 
-    with pytest.raises(TypeError):
-        throwing_task.delay()
+@retryable_task()
+@count_retries
+def throwing_task_programming_error():
+    raise ProgrammingError
+
+
+@retryable_task()
+@count_retries
+def throwing_task_unicode_error():
+    raise UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte")
+
+
+@retryable_task()
+@count_retries
+def throwing_task_validation_error():
+    raise jsonschema.ValidationError("Invalid schema")
+
+
+@retryable_task()
+@count_retries
+def throwing_task_zlib_error():
+    raise zlib.error("zlib compression error")
+
+
+@pytest.mark.parametrize(
+    "task_func, expected_exc",
+    [
+        (throwing_task_type_error, TypeError),
+        (throwing_task_key_error, KeyError),
+        (throwing_task_value_error, ValueError),
+        (throwing_task_index_error, IndexError),
+        (throwing_task_integrity_error, IntegrityError),
+        (throwing_task_programming_error, ProgrammingError),
+        (throwing_task_unicode_error, UnicodeDecodeError),
+        (throwing_task_validation_error, jsonschema.ValidationError),
+        (throwing_task_zlib_error, zlib.error),
+    ],
+)
+def test_retryable_task_throws_non_retryable(task_func, expected_exc):
+    "Test celery immediately raises an error for non-retryable exceptions without retrying"
+    with pytest.raises(expected_exc):
+        task_func.delay()
     assert thread_data.retry_count == 0
 
 
@@ -54,12 +123,42 @@ def throwing_task_should_retry():
 
 
 def test_retryable_task_throws_retry():
-    "Test celery executes a task properly"
+    "Test celery retry behavior on retryable exception"
 
     with pytest.raises(Retry) as e:
         throwing_task_should_retry.delay()
     assert str(e.value) == "Retry in 10s: OperationalError()"
+    assert thread_data.retry_count == 0
 
-    # The task is only called once, the Retry() exception
-    # will signal to the worker that the task needs to be tried again later
-    assert thread_data.retry_count == 1
+
+@retryable_task()
+@count_retries
+def throwing_missing_push():
+    raise MissingPushError("No push found")
+
+
+@retryable_task()
+@count_retries
+def throwing_runtime_error():
+    raise RuntimeError("Generic runtime error")
+
+
+@patch("newrelic.agent.notice_error")
+def test_newrelic_notified_on_generic_retryable_exception(mock_notice_error):
+    "Test that New Relic is notified on generic retryable exceptions"
+    with pytest.raises(Retry):
+        throwing_runtime_error.delay()
+
+    assert mock_notice_error.call_count == 1
+    mock_notice_error.assert_called_with(attributes={"number_of_prior_retries": 0})
+    assert thread_data.retry_count == 0
+
+
+@patch("newrelic.agent.notice_error")
+def test_newrelic_not_notified_on_hide_during_retries(mock_notice_error):
+    "Test that New Relic is NOT notified on exceptions in HIDE_DURING_RETRIES"
+    with pytest.raises(Retry):
+        throwing_missing_push.delay()
+
+    assert mock_notice_error.call_count == 0
+    assert thread_data.retry_count == 0