From d353c484bb40d6378273bc6a3ea6924e9c8e1b36 Mon Sep 17 00:00:00 2001
From: andrei <b1naryth1ef@gmail.com>
Date: Sat, 1 Dec 2018 04:59:39 -0800
Subject: [PATCH] Clear voice states when a new session connects (fixes #100,
 closes #109)

Discord is inconsistent here and doesn't emit a VoiceStateUpdate for the
old session_id when a user connects to voice from a different device
(thus disconnecting the first device). To avoid corrupting our voice
state cache, we need to remove any of the users old sessions (they
should have at most one).

I included tests on this one because its esoteric and undocumented, so
there is a high chance I'll break it in the future.
---
 disco/state.py |  6 ++++++
 tests/state.py | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 tests/state.py

diff --git a/disco/state.py b/disco/state.py
index e110c72..cdcdfe8 100644
--- a/disco/state.py
+++ b/disco/state.py
@@ -248,7 +248,13 @@ class State(object):
         # New connection
         elif event.state.channel_id:
             if event.state.guild_id in self.guilds:
+                expired_voice_state = self.guilds[event.state.guild_id].voice_states.select_one(user_id=event.user_id)
+                if expired_voice_state:
+                    del self.guilds[event.state.guild_id].voice_states[expired_voice_state.session_id]
                 self.guilds[event.state.guild_id].voice_states[event.state.session_id] = event.state
+            expired_voice_state = self.voice_states.select_one(user_id=event.user_id)
+            if expired_voice_state:
+                del self.voice_states[expired_voice_state.session_id]
             self.voice_states[event.state.session_id] = event.state
 
     def on_guild_member_add(self, event):
diff --git a/tests/state.py b/tests/state.py
new file mode 100644
index 0000000..dfe96d2
--- /dev/null
+++ b/tests/state.py
@@ -0,0 +1,39 @@
+from disco.state import State, StateConfig
+from holster.emitter import Emitter
+from disco.gateway.events import VoiceStateUpdate
+
+
+class MockClient(object):
+    def __init__(self):
+        self.events = Emitter()
+
+
+def get_state(config=None):
+    return State(MockClient(), config or StateConfig())
+
+
+def test_state_remove_expired_voice_states_device_change():
+    state = get_state()
+
+    event = VoiceStateUpdate.create({
+        'session_id': 'a',
+        'guild_id': 1,
+        'channel_id': 1,
+        'user_id': 1,
+    }, None)
+    state.client.events.emit('VoiceStateUpdate', event)
+
+    assert len(state.voice_states) == 1
+    assert 'a' in state.voice_states
+
+    event = VoiceStateUpdate.create({
+        'session_id': 'b',
+        'guild_id': 1,
+        'channel_id': 1,
+        'user_id': 1,
+    }, None)
+    state.client.events.emit('VoiceStateUpdate', event)
+
+    assert len(state.voice_states) == 1
+    assert 'a' not in state.voice_states
+    assert 'b' in state.voice_states