From 114ece1848625ef547655a01ade1fc44558cf474 Mon Sep 17 00:00:00 2001
From: Pascal Vizeli <pascal.vizeli@syshack.ch>
Date: Tue, 8 Nov 2016 10:24:50 +0100
Subject: [PATCH] Fix possible sigterm / unittest / Fix all lazy test (#4297)

* replace weakref with a list

* add unittest

* fix lint

* fix handling

* fix unittest

* change code style

* fix lazy tests
---
 homeassistant/core.py   | 16 +++++++++++-----
 homeassistant/remote.py |  3 +--
 tests/test_core.py      | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index 9a3211d8472..eea48d536e3 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -14,7 +14,6 @@ import re
 import signal
 import sys
 import threading
-import weakref
 
 from types import MappingProxyType
 from typing import Optional, Any, Callable, List  # NOQA
@@ -110,7 +109,7 @@ class HomeAssistant(object):
         self.executor = ThreadPoolExecutor(max_workers=5)
         self.loop.set_default_executor(self.executor)
         self.loop.set_exception_handler(self._async_exception_handler)
-        self._pending_tasks = weakref.WeakSet()
+        self._pending_tasks = []
         self.bus = EventBus(self)
         self.services = ServiceRegistry(self.bus, self.async_add_job,
                                         self.loop)
@@ -218,7 +217,12 @@ class HomeAssistant(object):
 
         # if a task is sheduled
         if task is not None:
-            self._pending_tasks.add(task)
+            self._pending_tasks.append(task)
+
+        # cleanup
+        if len(self._pending_tasks) > 50:
+            self._pending_tasks = [sheduled for sheduled in self._pending_tasks
+                                   if not sheduled.done()]
 
     @callback
     def async_run_job(self, target: Callable[..., None], *args: Any) -> None:
@@ -254,13 +258,15 @@ class HomeAssistant(object):
         """Block till all pending work is done."""
         while True:
             # Wait for the pending tasks are down
-            pending = list(self._pending_tasks)
+            pending = [task for task in self._pending_tasks
+                       if not task.done()]
+            self._pending_tasks.clear()
             if len(pending) > 0:
                 yield from asyncio.wait(pending, loop=self.loop)
 
             # Verify the loop is empty
             ret = yield from self.loop.run_in_executor(None, self._loop_empty)
-            if ret:
+            if ret and not self._pending_tasks:
                 break
 
     def stop(self) -> None:
diff --git a/homeassistant/remote.py b/homeassistant/remote.py
index ae96a397826..ba69d3249a9 100644
--- a/homeassistant/remote.py
+++ b/homeassistant/remote.py
@@ -16,7 +16,6 @@ import logging
 import time
 import threading
 import urllib.parse
-import weakref
 
 from typing import Optional
 
@@ -128,7 +127,7 @@ class HomeAssistant(ha.HomeAssistant):
         self.executor = ThreadPoolExecutor(max_workers=5)
         self.loop.set_default_executor(self.executor)
         self.loop.set_exception_handler(self._async_exception_handler)
-        self._pending_tasks = weakref.WeakSet()
+        self._pending_tasks = []
 
         self.bus = EventBus(remote_api, self)
         self.services = ha.ServiceRegistry(self.bus, self.add_job, self.loop)
diff --git a/tests/test_core.py b/tests/test_core.py
index 51653500b6e..025617f97fe 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -118,6 +118,43 @@ class TestHomeAssistant(unittest.TestCase):
 
     #     self.assertEqual(1, len(calls))
 
+    def test_async_add_job_pending_taks_add(self):
+        """Add a coro to pending tasks."""
+        call_count = []
+
+        @asyncio.coroutine
+        def test_coro():
+            """Test Coro."""
+            call_count.append('call')
+
+        self.hass.add_job(test_coro())
+
+        assert len(self.hass._pending_tasks) == 1
+        self.hass.block_till_done()
+        assert len(call_count) == 1
+
+    def test_async_add_job_pending_taks_cleanup(self):
+        """Add a coro to pending tasks."""
+        call_count = []
+
+        @asyncio.coroutine
+        def test_coro():
+            """Test Coro."""
+            call_count.append('call')
+
+        for i in range(50):
+            self.hass.add_job(test_coro())
+
+        assert len(self.hass._pending_tasks) == 50
+        self.hass.block_till_done()
+        assert len(call_count) == 50
+
+        self.hass.add_job(test_coro())
+
+        assert len(self.hass._pending_tasks) == 1
+        self.hass.block_till_done()
+        assert len(call_count) == 51
+
 
 class TestEvent(unittest.TestCase):
     """A Test Event class."""