From 2f5d4116ec49fff67983cd64acf82660200816f6 Mon Sep 17 00:00:00 2001
From: Chen Kasirer <ckasirer@ethz.ch>
Date: Mon, 24 Jul 2023 18:18:39 +0200
Subject: [PATCH] fixed graph serialization in assembly

---
 CHANGELOG.md                                   |  1 +
 src/compas/datastructures/assembly/assembly.py |  4 ++--
 tests/compas/datastructures/test_assembly.py   | 17 ++++++++++++++++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 643bc6867b0..216e6fb1ef1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -116,6 +116,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Changed `compas.data.DataEncoder` and `compas.data.DataDecoder` to support `to_jsondata` and `from_jsondata`.
 * Moved all API level docstrings from the `__init__.py` to the correspoding `.rst` file in the docs.
 * Fixed `AttributeError` in Plotter's `PolylineArtist` and `SegementArtist`.
+* Fixed wrong key type when de-serializing `Graph` with integer keys leading to node not found.
 
 ### Removed
 
diff --git a/src/compas/datastructures/assembly/assembly.py b/src/compas/datastructures/assembly/assembly.py
index 2407b346999..af6d33da85d 100644
--- a/src/compas/datastructures/assembly/assembly.py
+++ b/src/compas/datastructures/assembly/assembly.py
@@ -60,14 +60,14 @@ def JSONSCHEMANAME(self):
     def data(self):
         data = {
             "attributes": self.attributes,
-            "graph": self.graph.data,
+            "graph": self.graph,
         }
         return data
 
     @data.setter
     def data(self, data):
         self.attributes.update(data["attributes"] or {})
-        self.graph.data = data["graph"]
+        self.graph = data["graph"]
         self._parts = {part.guid: part.key for part in self.parts()}
 
     # ==========================================================================
diff --git a/tests/compas/datastructures/test_assembly.py b/tests/compas/datastructures/test_assembly.py
index 66a94626dfe..01e5149acb3 100644
--- a/tests/compas/datastructures/test_assembly.py
+++ b/tests/compas/datastructures/test_assembly.py
@@ -1,5 +1,7 @@
 import pytest
 
+from compas.data import json_dumps
+from compas.data import json_loads
 from compas.datastructures import Assembly
 from compas.datastructures import AssemblyError
 from compas.datastructures import Part
@@ -69,9 +71,22 @@ def test_find_by_key():
     assert assembly.find_by_key("100") is None
 
 
-def test_find_by_key_after_deserialization():
+def test_find_by_key_after_from_data():
     assembly = Assembly()
     part = Part()
     assembly.add_part(part, key=2)
     assembly = Assembly.from_data(assembly.to_data())
     assert assembly.find_by_key(2) == part
+
+
+def test_find_by_key_after_deserialization():
+    assembly = Assembly()
+    part = Part(name="test_part")
+    assembly.add_part(part, key=2)
+    assembly = json_loads(json_dumps(assembly))
+
+    deserialized_part = assembly.find_by_key(2)
+    assert deserialized_part.name == part.name
+    assert deserialized_part.key == part.key
+    assert deserialized_part.guid == part.guid
+    assert deserialized_part.attributes == part.attributes