From cf3a8b60ffa00eb144d56a8fac73f3b38c5bf655 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@home-assistant.io>
Date: Tue, 19 Feb 2019 23:02:56 -0800
Subject: [PATCH] Prevent invalid context from crashing (#21231)

* Prevent invalid context from crashing

* Lint
---
 homeassistant/core.py |   5 +-
 tests/test_core.py    | 105 ++++++++++++++++++++++++------------------
 2 files changed, 64 insertions(+), 46 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index c181ad453f3..48ef4f46272 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -744,7 +744,10 @@ class State:
 
         context = json_dict.get('context')
         if context:
-            context = Context(**context)
+            context = Context(
+                id=context.get('id'),
+                user_id=context.get('user_id'),
+            )
 
         return cls(json_dict['entity_id'], json_dict['state'],
                    json_dict.get('attributes'), last_changed, last_updated,
diff --git a/tests/test_core.py b/tests/test_core.py
index 818c6e4c087..e2ed249f441 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -460,61 +460,76 @@ class TestEventBus(unittest.TestCase):
         assert len(coroutine_calls) == 1
 
 
-class TestState(unittest.TestCase):
-    """Test State methods."""
+def test_state_init():
+    """Test state.init."""
+    with pytest.raises(InvalidEntityFormatError):
+        ha.State('invalid_entity_format', 'test_state')
 
-    def test_init(self):
-        """Test state.init."""
-        with pytest.raises(InvalidEntityFormatError):
-            ha.State('invalid_entity_format', 'test_state')
+    with pytest.raises(InvalidStateError):
+        ha.State('domain.long_state', 't' * 256)
 
-        with pytest.raises(InvalidStateError):
-            ha.State('domain.long_state', 't' * 256)
 
-    def test_domain(self):
-        """Test domain."""
-        state = ha.State('some_domain.hello', 'world')
-        assert 'some_domain' == state.domain
+def test_state_domain():
+    """Test domain."""
+    state = ha.State('some_domain.hello', 'world')
+    assert 'some_domain' == state.domain
 
-    def test_object_id(self):
-        """Test object ID."""
-        state = ha.State('domain.hello', 'world')
-        assert 'hello' == state.object_id
 
-    def test_name_if_no_friendly_name_attr(self):
-        """Test if there is no friendly name."""
-        state = ha.State('domain.hello_world', 'world')
-        assert 'hello world' == state.name
+def test_state_object_id():
+    """Test object ID."""
+    state = ha.State('domain.hello', 'world')
+    assert 'hello' == state.object_id
 
-    def test_name_if_friendly_name_attr(self):
-        """Test if there is a friendly name."""
-        name = 'Some Unique Name'
-        state = ha.State('domain.hello_world', 'world',
-                         {ATTR_FRIENDLY_NAME: name})
-        assert name == state.name
 
-    def test_dict_conversion(self):
-        """Test conversion of dict."""
-        state = ha.State('domain.hello', 'world', {'some': 'attr'})
-        assert state == ha.State.from_dict(state.as_dict())
+def test_state_name_if_no_friendly_name_attr():
+    """Test if there is no friendly name."""
+    state = ha.State('domain.hello_world', 'world')
+    assert 'hello world' == state.name
 
-    def test_dict_conversion_with_wrong_data(self):
-        """Test conversion with wrong data."""
-        assert ha.State.from_dict(None) is None
-        assert ha.State.from_dict({'state': 'yes'}) is None
-        assert ha.State.from_dict({'entity_id': 'yes'}) is None
 
-    def test_repr(self):
-        """Test state.repr."""
-        assert "<state happy.happy=on @ 1984-12-08T12:00:00+00:00>" == \
-            str(ha.State(
-                "happy.happy", "on",
-                last_changed=datetime(1984, 12, 8, 12, 0, 0)))
+def test_state_name_if_friendly_name_attr():
+    """Test if there is a friendly name."""
+    name = 'Some Unique Name'
+    state = ha.State('domain.hello_world', 'world',
+                     {ATTR_FRIENDLY_NAME: name})
+    assert name == state.name
 
-        assert "<state happy.happy=on; brightness=144 @ " \
-            "1984-12-08T12:00:00+00:00>" == \
-            str(ha.State("happy.happy", "on", {"brightness": 144},
-                         datetime(1984, 12, 8, 12, 0, 0)))
+
+def test_state_dict_conversion():
+    """Test conversion of dict."""
+    state = ha.State('domain.hello', 'world', {'some': 'attr'})
+    assert state == ha.State.from_dict(state.as_dict())
+
+
+def test_state_dict_conversion_with_wrong_data():
+    """Test conversion with wrong data."""
+    assert ha.State.from_dict(None) is None
+    assert ha.State.from_dict({'state': 'yes'}) is None
+    assert ha.State.from_dict({'entity_id': 'yes'}) is None
+    # Make sure invalid context data doesn't crash
+    wrong_context = ha.State.from_dict({
+        'entity_id': 'light.kitchen',
+        'state': 'on',
+        'context': {
+            'id': '123',
+            'non-existing': 'crash'
+        }
+    })
+    assert wrong_context is not None
+    assert wrong_context.context.id == '123'
+
+
+def test_state_repr():
+    """Test state.repr."""
+    assert "<state happy.happy=on @ 1984-12-08T12:00:00+00:00>" == \
+        str(ha.State(
+            "happy.happy", "on",
+            last_changed=datetime(1984, 12, 8, 12, 0, 0)))
+
+    assert "<state happy.happy=on; brightness=144 @ " \
+        "1984-12-08T12:00:00+00:00>" == \
+        str(ha.State("happy.happy", "on", {"brightness": 144},
+                     datetime(1984, 12, 8, 12, 0, 0)))
 
 
 class TestStateMachine(unittest.TestCase):