From 2e8afc1a4fc6e305c3ceff8a5847ee270ccd8c81 Mon Sep 17 00:00:00 2001
From: t92549 <80890692+t92549@users.noreply.github.com>
Date: Tue, 1 Mar 2022 11:44:19 +0000
Subject: [PATCH] gh-975: Fix gafferpy alpha 1

---
 .../src/gafferpy/gaffer_operations.py         | 22 ++++++++++++++++++-
 python-shell/src/test/test_connector.py       |  1 +
 .../src/test/test_connector_requests.py       |  1 +
 .../src/test/test_gaffer_operations.py        | 13 +++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/python-shell/src/gafferpy/gaffer_operations.py b/python-shell/src/gafferpy/gaffer_operations.py
index 0519f8ab6..ceab7815b 100755
--- a/python-shell/src/gafferpy/gaffer_operations.py
+++ b/python-shell/src/gafferpy/gaffer_operations.py
@@ -683,7 +683,7 @@ class GetTraits(Operation):
     CLASS = 'uk.gov.gchq.gaffer.store.operation.GetTraits'
 
     def __init__(self,
-                 current_traits,
+                 current_traits=True,
                  options=None):
         super().__init__(
             _class_name=self.CLASS, options=options)
@@ -696,6 +696,26 @@ def to_json(self):
         operation['currentTraits'] = self.current_traits
         return operation
 
+class HasTrait(Operation):
+    CLASS = 'uk.gov.gchq.gaffer.store.operation.HasTrait'
+
+    def __init__(self,
+                 trait,
+                 current_traits=True,
+                 options=None):
+        super().__init__(
+            _class_name=self.CLASS, options=options)
+
+        self.trait = trait
+        self.current_traits = current_traits
+
+    def to_json(self):
+        operation = super().to_json()
+
+        operation['trait'] = self.trait
+        operation['currentTraits'] = self.current_traits
+        return operation
+
 
 class AddElements(Operation):
     """
diff --git a/python-shell/src/test/test_connector.py b/python-shell/src/test/test_connector.py
index 4726794cb..95ca43b57 100755
--- a/python-shell/src/test/test_connector.py
+++ b/python-shell/src/test/test_connector.py
@@ -202,6 +202,7 @@ def test_execute_get(self):
             "uk.gov.gchq.gaffer.operation.impl.function.Filter",
             "uk.gov.gchq.gaffer.operation.impl.function.Transform",
             "uk.gov.gchq.gaffer.operation.impl.function.Aggregate",
+            "uk.gov.gchq.gaffer.store.operation.HasTrait",
             "uk.gov.gchq.gaffer.store.operation.GetTraits",
             "uk.gov.gchq.gaffer.mapstore.operation.CountAllElementsDefaultView",
             "uk.gov.gchq.gaffer.operation.export.graph.ExportToOtherAuthorisedGraph",
diff --git a/python-shell/src/test/test_connector_requests.py b/python-shell/src/test/test_connector_requests.py
index 218888287..9f48b4033 100755
--- a/python-shell/src/test/test_connector_requests.py
+++ b/python-shell/src/test/test_connector_requests.py
@@ -202,6 +202,7 @@ def test_execute_get(self):
             "uk.gov.gchq.gaffer.operation.impl.function.Filter",
             "uk.gov.gchq.gaffer.operation.impl.function.Transform",
             "uk.gov.gchq.gaffer.operation.impl.function.Aggregate",
+            "uk.gov.gchq.gaffer.store.operation.HasTrait",
             "uk.gov.gchq.gaffer.store.operation.GetTraits",
             "uk.gov.gchq.gaffer.mapstore.operation.CountAllElementsDefaultView",
             "uk.gov.gchq.gaffer.operation.export.graph.ExportToOtherAuthorisedGraph",
diff --git a/python-shell/src/test/test_gaffer_operations.py b/python-shell/src/test/test_gaffer_operations.py
index a7b8fab97..a44ce6629 100755
--- a/python-shell/src/test/test_gaffer_operations.py
+++ b/python-shell/src/test/test_gaffer_operations.py
@@ -4996,6 +4996,19 @@ class GafferOperationsTest(unittest.TestCase):
                 operation=g.GetElements()
             )
         ],
+        [
+            '''
+            {
+                "class" : "uk.gov.gchq.gaffer.store.operation.HasTrait",
+                "trait": "VISIBILITY",
+                "currentTraits" : true
+            }
+            ''',
+            g.HasTrait(
+                trait="VISIBILITY",
+                current_traits=True
+            )
+        ],
         [
             '''
             {