fixes #1835 (#1837)

* fixes #1835 * Add Electron __pow__ method * add tests to Electron __pow__ method * current version of test_electron_pow_method * fixing test_electron_pow after a change request * Empty commit --------- Co-authored-by: Sankalp Sanand <[email protected]>
AgnostiqHQ · Oct 19, 2023 · 9540267 · 9540267
1 parent 1995292
commit 9540267
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Documentation and test cases for database triggers.
+- Added the `__pow__` method to the `Electron` class
 
 ### Docs
 

diff --git a/covalent/_workflow/electron.py b/covalent/_workflow/electron.py
@@ -155,6 +155,7 @@ def get_op_function(
             "-": operator.sub,
             "*": operator.mul,
             "/": operator.truediv,
+            "**": operator.pow,
         }
 
         def rename(op1: Any, op: str, op2: Any) -> Callable:
@@ -243,6 +244,9 @@ def __truediv__(self, other):
     def __rtruediv__(self, other):
         return self.get_op_function(other, self, "/")
 
+    def __pow__(self, other):
+        return self.get_op_function(self, other, "**")
+
     def __int__(self):
         return int()
 

diff --git a/tests/covalent_tests/workflow/electron_test.py b/tests/covalent_tests/workflow/electron_test.py
@@ -17,7 +17,7 @@
 """Unit tests for electron"""
 
 import json
-from unittest.mock import MagicMock
+from unittest.mock import ANY, MagicMock
 
 import pytest
 
@@ -655,3 +655,22 @@ def workflow(x):
     assert (
         workflow.transport_graph.get_node_value(0, "status") == RESULT_STATUS.PENDING_REPLACEMENT
     )
+
+
+def test_electron_pow_method(mocker):
+    mock_electron_get_op_function = mocker.patch.object(
+        Electron, "get_op_function", return_value=Electron
+    )
+
+    @ct.electron
+    def g(x):
+        return 42 * x
+
+    @ct.lattice
+    def workflow(x):
+        res = g(x)
+        return res**2
+
+    workflow.build_graph(2)
+
+    mock_electron_get_op_function.assert_called_with(ANY, 2, "**")