diff --git a/src/main/native/org_bblfsh_client_v2_libuast_Libuast.cc b/src/main/native/org_bblfsh_client_v2_libuast_Libuast.cc
index e34b328..46f820c 100644
--- a/src/main/native/org_bblfsh_client_v2_libuast_Libuast.cc
+++ b/src/main/native/org_bblfsh_client_v2_libuast_Libuast.cc
@@ -50,6 +50,28 @@ jobject asJvmBuffer(uast::Buffer buf) {
   return env->NewDirectByteBuffer(buf.ptr, buf.size);
 }
 
+// Checks if a given object is of Node/JNode class
+bool isContext(jobject obj, JNIEnv *env) {
+  if (!obj) return false;
+
+  jclass ctxCls = env->FindClass(CLS_CTX);
+  checkJvmException("failed to find class " + std::string(CLS_CTX));
+
+  return env->IsInstanceOf(obj, ctxCls);
+}
+
+bool assertNotContext(jobject obj) {
+  JNIEnv *env = getJNIEnv();
+  if (isContext(obj, env)) {
+    auto reCls = env->FindClass(CLS_RE);
+    checkJvmException("failed to find class " + std::string(CLS_RE));
+
+    env->ThrowNew(reCls, "cannot use UAST Context as a Node");
+    return false;
+  }
+  return true;
+}
+
 // ==========================================
 // External UAST Context (managed by libuast)
 // ==========================================
@@ -76,7 +98,9 @@ class ContextExt {
     checkJvmException("failed to find class " + std::string(CLS_NODE));
 
     if (!env->IsInstanceOf(obj, cls)) {
-      const char *err = "ContextExt.toHandle() called not on Node type";
+      auto err = std::string("ContextExt.toHandle() called not on")
+                     .append(CLS_NODE)
+                     .append(" type");
       ctx->SetError(err);
       return 0;
     }
@@ -103,7 +127,7 @@ class ContextExt {
   // Encode serializes the external UAST.
   // Borrows the reference.
   jobject Encode(jobject node, UastFormat format) {
-    // if (!assertNotContext(node)) return nullptr;
+    if (!assertNotContext(node)) return nullptr;
 
     uast::Buffer data = ctx->Encode(toHandle(node), format);
     return asJvmBuffer(data);
@@ -358,7 +382,7 @@ class Interface : public uast::NodeCreator<Node *> {
   // toJ returns a JVM object associated with a node.
   // Returns a new reference.
   jobject toJ(Node *node) {
-    if (node == nullptr) return nullptr;
+    if (!node) return nullptr;
     jobject obj = getJNIEnv()->NewGlobalRef(node->obj);
     return obj;
   }
@@ -420,7 +444,7 @@ class Context {
   // toJ returns a JVM object associated with a node.
   // Returns a new reference.
   jobject toJ(Node *node) {
-    if (node == nullptr) return nullptr;
+    if (!node) return nullptr;
     return iface->toJ(node);
   }
   // toNode returns a node associated with a JVM object.
@@ -452,7 +476,7 @@ class Context {
   // Encode serializes UAST.
   // Creates a new reference.
   jobject Encode(jobject node, UastFormat format) {
-    // if (!assertNotContext(node)) return nullptr;
+    if (!assertNotContext(node)) return nullptr;
 
     Node *n = toNode(node);
     uast::Buffer data = ctx->Encode(n, format);
diff --git a/src/main/scala/org/bblfsh/client/v2/ContextExt.scala b/src/main/scala/org/bblfsh/client/v2/ContextExt.scala
index 96aca80..d3185f2 100644
--- a/src/main/scala/org/bblfsh/client/v2/ContextExt.scala
+++ b/src/main/scala/org/bblfsh/client/v2/ContextExt.scala
@@ -3,9 +3,9 @@ package org.bblfsh.client.v2
 import java.nio.ByteBuffer
 
 /**
-  * pyuast.ContextExt
-  *
   * Represents Go-side results of Libuast.decode()
+  *
+  * This is equivalent of pyuast.ContextExt API
   */
 case class ContextExt(nativeContext: Long) {
     @native def root(): Node
@@ -21,9 +21,9 @@ case class ContextExt(nativeContext: Long) {
 
 
 /**
-  * pyuast.Context
+  * Represents JVM-side constructed tree
   *
-  * Represents JVM-side constructed tree.
+  * This is equivalent of pyuast.Context API
   */
 case class Context(nativeContext: Long) {
     @native def root(): JNode
diff --git a/src/test/scala/org/bblfsh/client/v2/BblfshClientLoadTest.scala b/src/test/scala/org/bblfsh/client/v2/BblfshClientLoadTest.scala
index 26b1f41..c91eb4d 100644
--- a/src/test/scala/org/bblfsh/client/v2/BblfshClientLoadTest.scala
+++ b/src/test/scala/org/bblfsh/client/v2/BblfshClientLoadTest.scala
@@ -53,23 +53,24 @@ class BblfshClientLoadTest extends BblfshClientBaseTest {
   }
 
   "Decoding, loading & encoding to different context" should "produce the same results" in {
+    // decode -> load -> encode, and compare bytes
     val uast = resp.uast.decode()
     val rootNode: Node = uast.root()
 
     println(s"Loading $rootNode")
     val root = rootNode.load()
-
     val ctx = Context()
     val data = ctx.encode(root)
 
+    data should equal (resp.uast.asReadOnlyByteBuffer())
+
+    // decode -> load -> encoded -> decode -> load, and compare trees
     val uast2 = BblfshClient.decode(data)
     val rootNode2: Node = uast2.root()
     println(s"Loading $rootNode2")
     val root2 = rootNode2.load()
-    root2 should equal (root)
 
-    data should equal (resp.uast.asReadOnlyByteBuffer())
-//    data.compareTo(resp.uast.asReadOnlyByteBuffer()) shouldBe 0
+    root2 should equal (root)
   }
 
 }