Move tag and direct state copying into separate function

synapse/handlers/room_member.py

@@ -223,9 +223,36 @@ class RoomMemberHandler(object):
 
             # Copy over direct message status and room tags if this is a join
             # on an upgraded room
+            self.copy_room_tags_and_direct_to_room(old_room_id, room_id, user_id)
+        elif event.membership == Membership.LEAVE:
+            if prev_member_event_id:
+                prev_member_event = yield self.store.get_event(prev_member_event_id)
+                if prev_member_event.membership == Membership.JOIN:
+                    yield self._user_left_room(target, room_id)
 
+        defer.returnValue(event)
+
+    @defer.inlineCallbacks
+    def copy_room_tags_and_direct_to_room(
+        self,
+        old_room_id,
+        new_room_id,
+        user_id,
+    ):
+        """Copies the tags and direct room state from one room to another.
+
+        Args:
+            old_room_id (str)
+            new_room_id (str)
+            user_id (str)
+
+        Returns:
+            Deferred|None
+        """
         # Check if this is an upgraded room
-            state_ids = yield self.store.get_current_state_ids(room_id)
+        state_ids = yield self.store.get_filtered_current_state_ids(
+            new_room_id, StateFilter.from_types(((EventTypes.Create, ""))),
+        )
         create_id = state_ids.get((EventTypes.Create, ""))
         if not create_id:
             return
@@ -248,9 +275,9 @@ class RoomMemberHandler(object):
 
                 # Check which key this room is under
                 for key, room_id_list in direct_rooms.items():
-                        if old_room_id in room_id_list and room_id not in room_id_list:
+                    if old_room_id in room_id_list and new_room_id not in room_id_list:
                         # Add new room_id to this key
-                            direct_rooms[key].append(room_id)
+                        direct_rooms[key].append(new_room_id)
 
                         # Save back to user's m.direct account data
                         yield self.store.add_account_data_for_user(
@@ -266,13 +293,6 @@ class RoomMemberHandler(object):
                     yield self.store.add_tag_to_room(
                         user_id, room_id, tag, tag_content
                     )
-        elif event.membership == Membership.LEAVE:
-            if prev_member_event_id:
-                prev_member_event = yield self.store.get_event(prev_member_event_id)
-                if prev_member_event.membership == Membership.JOIN:
-                    yield self._user_left_room(target, room_id)
-
-        defer.returnValue(event)
 
     @defer.inlineCallbacks
     def update_membership(