Merge pull request #1823 from yahoo/geirst/impl-document-db-reference-resolver

Implement document db reference resolver

14 files changed, 338 insertions, 114 deletions

diff --git a/searchcore/CMakeLists.txt b/searchcore/CMakeLists.txt
index e98e1af15e6..882062baa64 100644
--- a/searchcore/CMakeLists.txt
+++ b/searchcore/CMakeLists.txt
@@ -119,6 +119,7 @@ vespa_define_module(
     src/tests/proton/proton
     src/tests/proton/reference/gid_to_lid_change_listener
     src/tests/proton/reference/gid_to_lid_mapper
+    src/tests/proton/reference/document_db_reference_resolver
     src/tests/proton/reference/document_db_referent_registry
     src/tests/proton/reprocessing/attribute_reprocessing_initializer
     src/tests/proton/reprocessing/document_reprocessing_handler
diff --git a/searchcore/src/tests/proton/common/cachedselect_test.cpp b/searchcore/src/tests/proton/common/cachedselect_test.cpp
index 9a3f08cc843..6d920a6bf27 100644
--- a/searchcore/src/tests/proton/common/cachedselect_test.cpp
+++ b/searchcore/src/tests/proton/common/cachedselect_test.cpp
@@ -15,6 +15,7 @@
 #include <vespa/searchlib/attribute/singlenumericenumattribute.hpp>
 #include <vespa/searchlib/attribute/singlenumericpostattribute.hpp>
 #include <vespa/searchlib/attribute/attributevector.hpp>
+#include <vespa/searchlib/test/mock_attribute_manager.h>
 #include <vespa/document/select/parser.h>
 #include <vespa/document/select/cloningvisitor.h>
 #include <vespa/document/base/documentid.h>
@@ -58,6 +59,7 @@ using search::SingleValueNumericPostingAttribute;
 using search::IntegerAttribute;
 using search::IntegerAttributeTemplate;
 using search::attribute::IAttributeContext;
+using search::attribute::test::MockAttributeManager;
 using namespace search::index;
 
 typedef Node::UP NodeUP;
@@ -206,69 +208,18 @@ public:
 };
 
 
-
-class MyAttributeManager : public search::IAttributeManager
+class MyAttributeManager : public MockAttributeManager
 {
 public:
-    typedef std::map<string, AttributeVector::SP> AttributeMap;
-
-    AttributeMap _attributes;
-
-    AttributeVector::SP
-    findAttribute(const vespalib::string &name) const
-    {
-        AttributeMap::const_iterator itr = _attributes.find(name);
-        if (itr != _attributes.end()) {
-            return itr->second;
-        }
-        return AttributeVector::SP();
-    }
-
-    virtual
-    AttributeGuard::UP
-    getAttribute(const string &name) const
-    {
-        AttributeVector::SP attr = findAttribute(name);
-        return AttributeGuard::UP(new AttributeGuard(attr));
-    }
-
-    virtual AttributeGuard::UP
-    getAttributeStableEnum(const string & name) const
-    {
-        AttributeVector::SP attr = findAttribute(name);
-        return AttributeGuard::UP(new AttributeEnumGuard(attr));
-    }
-
-    virtual void
-    getAttributeList(std::vector<AttributeGuard> & list) const
-    {
-        list.reserve(_attributes.size());
-        for (AttributeMap::const_iterator itr = _attributes.begin();
-             itr != _attributes.end();
-             ++itr) {
-            list.push_back(AttributeGuard(itr->second));
-        }
-    }
-
-    virtual IAttributeContext::UP
-    createContext() const
-    {
-        return IAttributeContext::UP(new AttributeContext(*this));
-    }
-
-    MyAttributeManager()
-        : _attributes()
-    {
-    }
-
-    void
-    addAttribute(const string &name)
-    {
-        if (findAttribute(name).get() != NULL)
+    void addAttribute(const vespalib::string &name) {
+        if (findAttribute(name).get() != NULL) {
             return;
+        }
         AttributeVector::SP av(new MyIntAv(name));
-        av->addReservedDoc();
-        _attributes[name] = av;
+        MockAttributeManager::addAttribute(name, av);
+    }
+    MyIntAv *getAsMyIntAttribute(const vespalib::string &name) const {
+        return (dynamic_cast<MyIntAv *>(findAttribute(name).get()));
     }
 };
 
@@ -320,7 +271,8 @@ MyDB::addDoc(uint32_t lid,
 
     _docIdToLid[docId] = lid;
     _lidToDocSP[lid] = Document::SP(doc.release());
-    AttributeVector &av(*_amgr.findAttribute("aa"));
+    AttributeGuard::UP guard = _amgr.getAttribute("aa");
+    AttributeVector &av = guard->get();
     if (lid >= av.getNumDocs()) {
         AttributeVector::DocId checkDocId(0u);
         ASSERT_TRUE(av.addDoc(checkDocId));
@@ -645,7 +597,7 @@ TEST_F("Test that basic select works", TestFixture)
     TEST_DO(checkSelect(cs, 3u, Result::Invalid));
     TEST_DO(checkSelect(cs, 4u, Result::Invalid));
 
-    MyIntAv *v(dynamic_cast<MyIntAv *>(f._amgr.findAttribute("aa").get()));
+    MyIntAv *v = f._amgr.getAsMyIntAttribute("aa");
     EXPECT_TRUE(v != NULL);
     EXPECT_EQUAL(6u, v->getGets());
 }
diff --git a/searchcore/src/tests/proton/document_iterator/document_iterator_test.cpp b/searchcore/src/tests/proton/document_iterator/document_iterator_test.cpp
index e42ed61dde9..b4bf103bf65 100644
--- a/searchcore/src/tests/proton/document_iterator/document_iterator_test.cpp
+++ b/searchcore/src/tests/proton/document_iterator/document_iterator_test.cpp
@@ -1,18 +1,19 @@
 // Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
-#include <vespa/document/fieldvalue/fieldvalues.h>
 #include <vespa/document/fieldset/fieldsets.h>
+#include <vespa/document/fieldvalue/fieldvalues.h>
 #include <vespa/persistence/spi/bucket.h>
 #include <vespa/persistence/spi/docentry.h>
 #include <vespa/persistence/spi/result.h>
+#include <vespa/searchcommon/common/schema.h>
+#include <vespa/searchcore/proton/common/attrupdate.h>
 #include <vespa/searchcore/proton/persistenceengine/document_iterator.h>
-#include <vespa/vespalib/testkit/test_kit.h>
+#include <vespa/searchcore/proton/server/commit_and_wait_document_retriever.h>
 #include <vespa/searchlib/attribute/attributecontext.h>
-#include <vespa/searchcore/proton/common/attrupdate.h>
 #include <vespa/searchlib/attribute/attributefactory.h>
-#include <vespa/searchcore/proton/server/commit_and_wait_document_retriever.h>
+#include <vespa/searchlib/test/mock_attribute_manager.h>
 #include <vespa/vespalib/objects/nbostream.h>
-#include <vespa/searchcommon/common/schema.h>
+#include <vespa/vespalib/testkit/test_kit.h>
 
 using document::DocumentType;
 using document::Field;
@@ -35,6 +36,7 @@ using search::attribute::BasicType;
 using search::attribute::CollectionType;
 using search::attribute::Config;
 using search::attribute::IAttributeContext;
+using search::attribute::test::MockAttributeManager;
 using namespace search::index;
 using storage::spi::Timestamp;
 using storage::spi::Bucket;
@@ -168,56 +170,9 @@ struct VisitRecordingUnitDR : UnitDR {
     }
 };
 
-class MyAttributeManager : public search::IAttributeManager
-{
-public:
-    typedef std::map<string, AttributeVector::SP> AttributeMap;
-
-    AttributeMap _attributes;
-
-    AttributeVector::SP
-    findAttribute(const vespalib::string &name) const {
-        AttributeMap::const_iterator itr = _attributes.find(name);
-        if (itr != _attributes.end()) {
-            return itr->second;
-        }
-        return AttributeVector::SP();
-    }
-
-    AttributeGuard::UP getAttribute(const string &name) const override {
-        AttributeVector::SP attr = findAttribute(name);
-        return AttributeGuard::UP(new AttributeGuard(attr));
-    }
-
-    AttributeGuard::UP getAttributeStableEnum(const string & name) const override {
-        AttributeVector::SP attr = findAttribute(name);
-        return AttributeGuard::UP(new AttributeEnumGuard(attr));
-    }
-
-    void getAttributeList(std::vector<AttributeGuard> & list) const override {
-        list.reserve(_attributes.size());
-        for (AttributeMap::const_iterator itr = _attributes.begin();
-             itr != _attributes.end();
-             ++itr) {
-            list.push_back(AttributeGuard(itr->second));
-        }
-    }
-
-    IAttributeContext::UP createContext() const override {
-        return IAttributeContext::UP(new AttributeContext(*this));
-    }
-
-    MyAttributeManager() : _attributes() { }
-
-    void addAttribute(const string &name, const AttributeVector::SP &av) {
-        av->addReservedDoc();
-        _attributes[name] = av;
-    }
-};
-
 struct AttrUnitDR : public UnitDR
 {
-    MyAttributeManager _amgr;
+    MockAttributeManager _amgr;
     Schema _schema;
     AttributeVector::SP _aa;
     AttributeVector::SP _dd;
diff --git a/searchcore/src/tests/proton/reference/document_db_reference_resolver/CMakeLists.txt b/searchcore/src/tests/proton/reference/document_db_reference_resolver/CMakeLists.txt
new file mode 100644
index 00000000000..1a2aa573ad1
--- /dev/null
+++ b/searchcore/src/tests/proton/reference/document_db_reference_resolver/CMakeLists.txt
@@ -0,0 +1,8 @@
+# Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(searchcore_document_db_reference_resolver_test_app TEST
+    SOURCES
+    document_db_reference_resolver_test.cpp
+    DEPENDS
+    searchcore_reference
+)
+vespa_add_test(NAME searchcore_document_db_reference_resolver_test_app COMMAND searchcore_document_db_reference_resolver_test_app)
diff --git a/searchcore/src/tests/proton/reference/document_db_reference_resolver/DESC b/searchcore/src/tests/proton/reference/document_db_reference_resolver/DESC
new file mode 100644
index 00000000000..3a8fe966966
--- /dev/null
+++ b/searchcore/src/tests/proton/reference/document_db_reference_resolver/DESC
@@ -0,0 +1 @@
+document_db_reference_resolver test. Take a look at document_db_reference_resolver_test.cpp for details.
diff --git a/searchcore/src/tests/proton/reference/document_db_reference_resolver/FILES b/searchcore/src/tests/proton/reference/document_db_reference_resolver/FILES
new file mode 100644
index 00000000000..15233d66c94
--- /dev/null
+++ b/searchcore/src/tests/proton/reference/document_db_reference_resolver/FILES
@@ -0,0 +1 @@
+document_db_reference_resolver_test.cpp
diff --git a/searchcore/src/tests/proton/reference/document_db_reference_resolver/document_db_reference_resolver_test.cpp b/searchcore/src/tests/proton/reference/document_db_reference_resolver/document_db_reference_resolver_test.cpp
new file mode 100644
index 00000000000..c8158d7d8c5
--- /dev/null
+++ b/searchcore/src/tests/proton/reference/document_db_reference_resolver/document_db_reference_resolver_test.cpp
@@ -0,0 +1,130 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#include <vespa/vespalib/testkit/testapp.h>
+
+#include <vespa/document/datatype/documenttype.h>
+#include <vespa/document/datatype/referencedatatype.h>
+#include <vespa/log/log.h>
+#include <vespa/searchcore/proton/reference/document_db_reference_resolver.h>
+#include <vespa/searchcore/proton/reference/i_document_db_referent.h>
+#include <vespa/searchcore/proton/reference/i_document_db_referent_registry.h>
+#include <vespa/searchlib/attribute/attributefactory.h>
+#include <vespa/searchlib/attribute/reference_attribute.h>
+#include <vespa/searchlib/common/i_gid_to_lid_mapper.h>
+#include <vespa/searchlib/common/i_gid_to_lid_mapper_factory.h>
+#include <vespa/searchlib/test/mock_attribute_manager.h>
+
+LOG_SETUP("document_db_reference_resolver_test");
+
+using namespace document;
+using namespace proton;
+using namespace search::attribute;
+using namespace search;
+
+struct MyGidToLidMapperFactory : public IGidToLidMapperFactory {
+    using SP = std::shared_ptr<MyGidToLidMapperFactory>;
+    virtual std::unique_ptr<IGidToLidMapper> getMapper() const override {
+        return std::unique_ptr<IGidToLidMapper>();
+    }
+};
+
+struct MyDocumentDBReferent : public IDocumentDBReferent {
+    MyGidToLidMapperFactory::SP factory;
+    MyDocumentDBReferent(MyGidToLidMapperFactory::SP factory_) : factory(factory_) {}
+    virtual AttributeVector::SP getAttribute(vespalib::stringref) override {
+        return AttributeVector::SP();
+    }
+    virtual IGidToLidMapperFactory::SP getGidToLidMapperFactory() override {
+        return factory;
+    }
+};
+
+struct MyReferentRegistry : public IDocumentDBReferentRegistry {
+    using ReferentMap = std::map<vespalib::string, IDocumentDBReferent::SP>;
+    ReferentMap map;
+    virtual IDocumentDBReferent::SP get(vespalib::stringref name) const override {
+        auto itr = map.find(name);
+        if (itr != map.end()) {
+            return itr->second;
+        }
+        return IDocumentDBReferent::SP();
+    }
+    virtual void add(vespalib::stringref name, IDocumentDBReferent::SP referent) override {
+        map[name] = referent;
+    }
+    virtual void remove(vespalib::stringref) override {}
+};
+
+struct MyAttributeManager : public test::MockAttributeManager {
+    void addIntAttribute(const vespalib::string &name) {
+        addAttribute(name, AttributeFactory::createAttribute(name, Config(BasicType::INT32)));
+    }
+    void addReferenceAttribute(const vespalib::string &name) {
+        addAttribute(name, std::make_shared<ReferenceAttribute>(name, Config(BasicType::REFERENCE)));
+    }
+    const ReferenceAttribute *getReferenceAttribute(const vespalib::string &name) const {
+        AttributeGuard::UP guard = getAttribute(name);
+        const ReferenceAttribute *result = dynamic_cast<const ReferenceAttribute *>(&guard->get());
+        ASSERT_TRUE(result != nullptr);
+        return result;
+    }
+};
+
+struct DocumentModel {
+    DocumentType parentDocType;
+    ReferenceDataType refDataType;
+    DocumentType childDocType;
+    DocumentModel()
+        : parentDocType("parent"),
+          refDataType(parentDocType, 1234),
+          childDocType("child")
+    {
+        initChildDocType();
+    }
+    void initChildDocType() {
+        childDocType.addField(Field("ref", refDataType, true));
+        childDocType.addField(Field("other_ref", refDataType, true));
+    }
+};
+
+struct Fixture {
+    MyGidToLidMapperFactory::SP factory;
+    MyDocumentDBReferent::SP parentReferent;
+    MyReferentRegistry registry;
+    MyAttributeManager attrMgr;
+    DocumentModel docModel;
+    DocumentDBReferenceResolver resolver;
+    Fixture() :
+        factory(std::make_shared<MyGidToLidMapperFactory>()),
+        parentReferent(std::make_shared<MyDocumentDBReferent>(factory)),
+        registry(),
+        attrMgr(),
+        docModel(),
+        resolver(registry, attrMgr, docModel.childDocType)
+    {
+        registry.add("parent", parentReferent);
+        populateAttributeManager();
+    }
+    void populateAttributeManager() {
+        attrMgr.addReferenceAttribute("ref");
+        attrMgr.addReferenceAttribute("other_ref");
+        attrMgr.addIntAttribute("int_attr");
+    }
+    void resolve() {
+        resolver.resolve();
+    }
+    const IGidToLidMapperFactory *getMapperFactoryPtr(const vespalib::string &attrName) {
+        return attrMgr.getReferenceAttribute(attrName)->getGidToLidMapperFactory().get();
+    }
+};
+
+TEST_F("require that reference attributes are connected to gid mapper", Fixture)
+{
+    f.resolve();
+    EXPECT_EQUAL(f.factory.get(), f.getMapperFactoryPtr("ref"));
+    EXPECT_EQUAL(f.factory.get(), f.getMapperFactoryPtr("other_ref"));
+}
+
+TEST_MAIN()
+{
+    TEST_RUN_ALL();
+}
diff --git a/searchcore/src/vespa/searchcore/proton/reference/CMakeLists.txt b/searchcore/src/vespa/searchcore/proton/reference/CMakeLists.txt
index 405e78ab2d2..94d7cec68ea 100644
--- a/searchcore/src/vespa/searchcore/proton/reference/CMakeLists.txt
+++ b/searchcore/src/vespa/searchcore/proton/reference/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 vespa_add_library(searchcore_reference STATIC
     SOURCES
+    document_db_reference_resolver.cpp
     document_db_referent.cpp
     document_db_referent_registry.cpp
     gid_to_lid_change_listener.cpp
diff --git a/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.cpp b/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.cpp
new file mode 100644
index 00000000000..244bfcc6218
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.cpp
@@ -0,0 +1,78 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/fastos/fastos.h>
+#include "document_db_reference_resolver.h"
+#include "i_document_db_referent_registry.h"
+#include "i_document_db_referent.h"
+#include <vespa/document/datatype/documenttype.h>
+#include <vespa/document/repo/documenttyperepo.h>
+#include <vespa/searchlib/attribute/attributevector.h>
+#include <vespa/searchlib/attribute/iattributemanager.h>
+#include <vespa/searchlib/attribute/reference_attribute.h>
+#include <vespa/document/datatype/referencedatatype.h>
+
+using document::DataType;
+using document::DocumentType;
+using document::DocumentTypeRepo;
+using document::ReferenceDataType;
+using search::attribute::BasicType;
+using search::attribute::ReferenceAttribute;
+using search::AttributeGuard;
+using search::AttributeVector;
+using search::IAttributeManager;
+
+namespace proton {
+
+namespace {
+
+vespalib::string
+getTargetDocTypeName(const vespalib::string &attrName,
+                     const DocumentType &thisDocType)
+{
+    const DataType &attrType = thisDocType.getField(attrName).getDataType();
+    const ReferenceDataType *refType = dynamic_cast<const ReferenceDataType *>(&attrType);
+    assert(refType != nullptr);
+    return refType->getTargetType().getName();
+}
+
+ReferenceAttribute &
+asReferenceAttribute(AttributeVector &attr)
+{
+    ReferenceAttribute *result = dynamic_cast<ReferenceAttribute *>(&attr);
+    assert(result != nullptr);
+    return *result;
+}
+
+}
+
+void
+DocumentDBReferenceResolver::connectReferenceAttributesToGidMapper()
+{
+    std::vector<AttributeGuard> attributeList;
+    _attrMgr.getAttributeList(attributeList);
+    for (auto &guard : attributeList) {
+        AttributeVector &attr = guard.get();
+        if (attr.getBasicType() == BasicType::REFERENCE) {
+            IDocumentDBReferent::SP targetDB = _registry.get(getTargetDocTypeName(attr.getName(),
+                                                                                  _thisDocType));
+            asReferenceAttribute(attr).setGidToLidMapperFactory(targetDB->getGidToLidMapperFactory());
+        }
+    }
+}
+
+DocumentDBReferenceResolver::DocumentDBReferenceResolver(const IDocumentDBReferentRegistry &registry,
+                                                         const IAttributeManager &attrMgr,
+                                                         const DocumentType &thisDocType)
+    : _registry(registry),
+      _attrMgr(attrMgr),
+      _thisDocType(thisDocType)
+{
+}
+
+void
+DocumentDBReferenceResolver::resolve()
+{
+    connectReferenceAttributesToGidMapper();
+}
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.h b/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.h
new file mode 100644
index 00000000000..efb6ef6a76c
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/reference/document_db_reference_resolver.h
@@ -0,0 +1,34 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#pragma once
+
+namespace document {
+class DocumentType;
+class DocumentTypeRepo;
+}
+namespace search { class IAttributeManager; }
+
+namespace proton {
+
+class IDocumentDBReferentRegistry;
+
+/**
+ * Class that for a given document db resolves all references to parent document dbs:
+ *   1) Connects all reference attributes to gid mappers of parent document dbs.
+ */
+class DocumentDBReferenceResolver {
+private:
+    const IDocumentDBReferentRegistry &_registry;
+    const search::IAttributeManager &_attrMgr;
+    const document::DocumentType &_thisDocType;
+
+    void connectReferenceAttributesToGidMapper();
+
+public:
+    DocumentDBReferenceResolver(const IDocumentDBReferentRegistry &registry,
+                                const search::IAttributeManager &attrMgr,
+                                const document::DocumentType &thisDocType);
+
+    void resolve();
+};
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/reference/i_document_db_referent.h b/searchcore/src/vespa/searchcore/proton/reference/i_document_db_referent.h
index 695539bf591..8ae4a28b185 100644
--- a/searchcore/src/vespa/searchcore/proton/reference/i_document_db_referent.h
+++ b/searchcore/src/vespa/searchcore/proton/reference/i_document_db_referent.h
@@ -20,6 +20,7 @@ namespace proton {
 class IDocumentDBReferent
 {
 public:
+    using SP = std::shared_ptr<IDocumentDBReferent>;
     virtual ~IDocumentDBReferent() { }
     virtual std::shared_ptr<search::AttributeVector> getAttribute(vespalib::stringref name) = 0;
     virtual std::shared_ptr<search::IGidToLidMapperFactory> getGidToLidMapperFactory() = 0;
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
index 1d7c53f83a1..ebff6c2bff6 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
@@ -74,6 +74,7 @@ public:
     void update(DocId doc, const GlobalId &gid);
     const Reference *getReference(DocId doc);
     void setGidToLidMapperFactory(std::shared_ptr<IGidToLidMapperFactory> gidToLidMapperFactory);
+    std::shared_ptr<IGidToLidMapperFactory> getGidToLidMapperFactory() const { return _gidToLidMapperFactory; }
     DocId getReferencedLid(DocId doc) const;
     void notifyGidToLidChange(const GlobalId &gid, DocId referencedLid);
 };
diff --git a/searchlib/src/vespa/searchlib/common/i_gid_to_lid_mapper_factory.h b/searchlib/src/vespa/searchlib/common/i_gid_to_lid_mapper_factory.h
index fadd0b74020..5f57ada4486 100644
--- a/searchlib/src/vespa/searchlib/common/i_gid_to_lid_mapper_factory.h
+++ b/searchlib/src/vespa/searchlib/common/i_gid_to_lid_mapper_factory.h
@@ -12,6 +12,7 @@ class IGidToLidMapper;
 class IGidToLidMapperFactory
 {
 public:
+    using SP = std::shared_ptr<IGidToLidMapperFactory>;
     virtual ~IGidToLidMapperFactory() { }
     virtual std::unique_ptr<IGidToLidMapper> getMapper() const = 0;
 };
diff --git a/searchlib/src/vespa/searchlib/test/mock_attribute_manager.h b/searchlib/src/vespa/searchlib/test/mock_attribute_manager.h
new file mode 100644
index 00000000000..bd5e0b0b8d2
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/test/mock_attribute_manager.h
@@ -0,0 +1,60 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#pragma once
+
+#include <vespa/searchlib/attribute/attributecontext.h>
+#include <vespa/searchlib/attribute/attributeguard.h>
+#include <vespa/searchlib/attribute/attributevector.h>
+#include <vespa/searchlib/attribute/iattributemanager.h>
+
+namespace search {
+namespace attribute {
+namespace test {
+
+class MockAttributeManager : public search::IAttributeManager {
+protected:
+    using AttributeMap = std::map<string, AttributeVector::SP>;
+
+    AttributeMap _attributes;
+
+    AttributeVector::SP
+    findAttribute(const vespalib::string &name) const {
+        AttributeMap::const_iterator itr = _attributes.find(name);
+        if (itr != _attributes.end()) {
+            return itr->second;
+        }
+        return AttributeVector::SP();
+    }
+
+public:
+    MockAttributeManager() : _attributes() {}
+
+    virtual AttributeGuard::UP getAttribute(const vespalib::string &name) const override {
+        AttributeVector::SP attr = findAttribute(name);
+        return AttributeGuard::UP(new AttributeGuard(attr));
+    }
+
+    virtual AttributeGuard::UP getAttributeStableEnum(const vespalib::string &name) const override {
+        AttributeVector::SP attr = findAttribute(name);
+        return AttributeGuard::UP(new AttributeEnumGuard(attr));
+    }
+
+    virtual void getAttributeList(std::vector<AttributeGuard> &list) const override {
+        list.reserve(_attributes.size());
+        for (const auto &attr : _attributes) {
+            list.push_back(AttributeGuard(attr.second));
+        }
+    }
+
+    virtual IAttributeContext::UP createContext() const override {
+        return IAttributeContext::UP(new AttributeContext(*this));
+    }
+
+    void addAttribute(const vespalib::string &name, const AttributeVector::SP &attr) {
+        attr->addReservedDoc();
+        _attributes[name] = attr;
+    }
+};
+
+}
+}
+}