# coding: utf-8

from qgis.core import (
    QgsVectorJoinInfo, QgsVectorLayer, QgsField, QgsFeature,
    QgsMapLayerRegistry, edit)
from qgis.PyQt import QVariant

# Create layer_1
layer_1 = QgsVectorLayer('Point?crs=epsg:21781', 'layer_1', 'memory')

with edit(layer_1):
    layer_1.addAttribute(QgsField('id', QVariant.Int))
    layer_1.addAttribute(QgsField('column_a', QVariant.String))
    layer_1.addAttribute(QgsField('column_b', QVariant.String))

    feat = QgsFeature()
    feat.setAttributes([1, 'A1', 'B1'])
    layer_1.addFeature(feat)

# add layer to the legend
QgsMapLayerRegistry.instance().addMapLayer(layer_1)

# Create layer_2
layer_2 = QgsVectorLayer('Point?crs=epsg:21781', 'layer_2', 'memory')

with edit(layer_2):
    layer_2.addAttribute(QgsField('id', QVariant.Int))

    feat = QgsFeature()
    feat.setAttributes([1])
    layer_2.addFeature(feat)

# add layer to the legend
QgsMapLayerRegistry.instance().addMapLayer(layer_2)

# Join the layers with the subset filds in the same order than the joined
# layer
join = QgsVectorJoinInfo()
join.joinLayerId = layer_1.id()
join.joinFieldName = 'id'
join.targetFieldName = 'id'
join.setJoinFieldNamesSubset(['column_a', 'column_b'])
layer_2.addJoin(join)

# Verify the result
assert layer_2.getFeatures().next().attribute(
        'layer_1_column_a') == 'A1', 'Join failed'

# Remove the join
layer_2.removeJoin(layer_1.id())

# Join the layers with the subset filds in a different order than the joined
# layer
join = QgsVectorJoinInfo()
join.joinLayerId = layer_1.id()
join.joinFieldName = 'id'
join.targetFieldName = 'id'
join.setJoinFieldNamesSubset(['column_b', 'column_a'])
layer_2.addJoin(join)

# Verify the result
assert layer_2.getFeatures().next().attribute(
        'layer_1_column_a') == 'A1', 'Join failed'