ec: Fix object comparison in the stack analyzer

Handle the case of comparing with different kinds of objects.

BUG=none
BRANCH=none
TEST=extra/stack_analyzer/stack_analyzer_unittest.py

Change-Id: I01056cd39e14d75442d4029b6c64d9843c49cf2a
Signed-off-by: Che-yu Wu <cheyuw@google.com>
Reviewed-on: https://chromium-review.googlesource.com/616367
Reviewed-by: Nicolas Boichat <drinkcat@chromium.org>

2 files changed, 44 insertions, 0 deletions

diff --git a/extra/stack_analyzer/stack_analyzer.py b/extra/stack_analyzer/stack_analyzer.py
index df7b7c8932..cc9543cc8c 100755
--- a/extra/stack_analyzer/stack_analyzer.py
+++ b/extra/stack_analyzer/stack_analyzer.py
@@ -76,6 +76,9 @@ class Task(object):
     Returns:
       True if equal, False if not.
     """
+    if not isinstance(other, Task):
+      return False
+
     return (self.name == other.name and
             self.routine_name == other.routine_name and
             self.stack_max_size == other.stack_max_size and
@@ -116,6 +119,9 @@ class Symbol(object):
     Returns:
       True if equal, False if not.
     """
+    if not isinstance(other, Symbol):
+      return False
+
     return (self.address == other.address and
             self.symtype == other.symtype and
             self.size == other.size and
@@ -156,6 +162,9 @@ class Callsite(object):
     Returns:
       True if equal, False if not.
     """
+    if not isinstance(other, Callsite):
+      return False
+
     if not (self.address == other.address and
             self.target == other.target and
             self.is_tail == other.is_tail):
@@ -217,6 +226,9 @@ class Function(object):
     Returns:
       True if equal, False if not.
     """
+    if not isinstance(other, Function):
+      return False
+
     # TODO(cheyuw): Don't compare SCC node attributes here.
     if not (self.address == other.address and
             self.name == other.name and
diff --git a/extra/stack_analyzer/stack_analyzer_unittest.py b/extra/stack_analyzer/stack_analyzer_unittest.py
index ff8ec840f3..390072fd8b 100755
--- a/extra/stack_analyzer/stack_analyzer_unittest.py
+++ b/extra/stack_analyzer/stack_analyzer_unittest.py
@@ -14,6 +14,38 @@ import unittest
 import stack_analyzer as sa
 
 
+class ObjectTest(unittest.TestCase):
+  """Tests for classes of basic objects."""
+
+  def testTask(self):
+    task_a = sa.Task('a', 'a_task', 1234)
+    task_b = sa.Task('b', 'b_task', 5678, 0x1000)
+    self.assertEqual(task_a, task_a)
+    self.assertNotEqual(task_a, task_b)
+    self.assertNotEqual(task_a, None)
+
+  def testSymbol(self):
+    symbol_a = sa.Symbol(0x1234, 'F', 32, 'a')
+    symbol_b = sa.Symbol(0x234, 'O', 42, 'b')
+    self.assertEqual(symbol_a, symbol_a)
+    self.assertNotEqual(symbol_a, symbol_b)
+    self.assertNotEqual(symbol_a, None)
+
+  def testCallsite(self):
+    callsite_a = sa.Callsite(0x1002, 0x3000, False)
+    callsite_b = sa.Callsite(0x1002, 0x3000, True)
+    self.assertEqual(callsite_a, callsite_a)
+    self.assertNotEqual(callsite_a, callsite_b)
+    self.assertNotEqual(callsite_a, None)
+
+  def testFunction(self):
+    func_a = sa.Function(0x100, 'a', 0, [])
+    func_b = sa.Function(0x200, 'b', 0, [])
+    self.assertEqual(func_a, func_a)
+    self.assertNotEqual(func_a, func_b)
+    self.assertNotEqual(func_a, None)
+
+
 class ArmAnalyzerTest(unittest.TestCase):
   """Tests for class ArmAnalyzer."""