IVGCVSW-2552 Add unit tests for the input and output layers for Visitor

Change-Id: I73ae9b43fa1631fdec423177d37fffc719c71f88
Signed-off-by: FrancisMurtagh <francis.murtagh@arm.com>

1 files changed, 61 insertions, 0 deletions

diff --git a/src/armnn/test/TestInputOutputLayerVisitor.hpp b/src/armnn/test/TestInputOutputLayerVisitor.hpp
new file mode 100644
index 0000000000..f67a65cd09
--- /dev/null
+++ b/src/armnn/test/TestInputOutputLayerVisitor.hpp
@@ -0,0 +1,61 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+#pragma once
+
+#include "TestLayerVisitor.hpp"
+#include <boost/test/unit_test.hpp>
+
+namespace armnn
+{
+
+void CheckLayerBindingId(LayerBindingId visitorId, LayerBindingId id)
+{
+    BOOST_CHECK_EQUAL(visitorId, id);
+}
+
+// Concrete TestLayerVisitor subclasses for layers taking LayerBindingId argument with overridden VisitLayer methods
+class TestInputLayerVisitor : public TestLayerVisitor
+{
+private:
+    LayerBindingId visitorId;
+
+public:
+    explicit TestInputLayerVisitor(LayerBindingId id, const char* name = nullptr)
+        : TestLayerVisitor(name)
+        , visitorId(id)
+    {};
+
+    void VisitInputLayer(const IConnectableLayer* layer,
+                         LayerBindingId id,
+                         const char* name = nullptr) override
+    {
+        CheckLayerPointer(layer);
+        CheckLayerBindingId(visitorId, id);
+        CheckLayerName(name);
+    };
+};
+
+class TestOutputLayerVisitor : public TestLayerVisitor
+{
+private:
+    LayerBindingId visitorId;
+
+public:
+    explicit TestOutputLayerVisitor(LayerBindingId id, const char* name = nullptr)
+        : TestLayerVisitor(name)
+        , visitorId(id)
+    {};
+
+    void VisitOutputLayer(const IConnectableLayer* layer,
+                          LayerBindingId id,
+                          const char* name = nullptr) override
+    {
+        CheckLayerPointer(layer);
+        CheckLayerBindingId(visitorId, id);
+        CheckLayerName(name);
+    };
+};
+
+} //namespace armnn