1# -*- coding: utf-8 -*-
2"""QGIS Unit tests for edit widgets.
3
4.. note:: This program is free software; you can redistribute it and/or modify
5it under the terms of the GNU General Public License as published by
6the Free Software Foundation; either version 2 of the License, or
7(at your option) any later version.
8"""
9__author__ = 'Matthias Kuhn'
10__date__ = '28/11/2015'
11__copyright__ = 'Copyright 2015, The QGIS Project'
12
13import qgis  # NOQA
14
15import os
16
17from qgis.core import (
18    QgsFeature,
19    QgsVectorLayer,
20    QgsProject,
21    QgsRelation,
22    QgsTransaction,
23    QgsFeatureRequest,
24    QgsVectorLayerTools,
25    QgsGeometry
26)
27
28from qgis.gui import (
29    QgsGui,
30    QgsRelationWidgetWrapper,
31    QgsAttributeEditorContext,
32    QgsMapCanvas,
33    QgsAdvancedDigitizingDockWidget
34)
35
36from qgis.PyQt.QtCore import QTimer
37from qgis.PyQt.QtWidgets import (
38    QToolButton,
39    QMessageBox,
40    QDialogButtonBox,
41    QTableView,
42    QDialog
43)
44from qgis.testing import start_app, unittest
45
46start_app()
47
48
49class TestQgsRelationEditWidget(unittest.TestCase):
50
51    @classmethod
52    def setUpClass(cls):
53        """
54        Setup the involved layers and relations for a n:m relation
55        :return:
56        """
57        cls.mapCanvas = QgsMapCanvas()
58        QgsGui.editorWidgetRegistry().initEditors(cls.mapCanvas)
59        cls.dbconn = 'service=\'qgis_test\''
60        if 'QGIS_PGTEST_DB' in os.environ:
61            cls.dbconn = os.environ['QGIS_PGTEST_DB']
62        # Create test layer
63        cls.vl_books = QgsVectorLayer(cls.dbconn + ' sslmode=disable key=\'pk\' table="qgis_test"."books" sql=', 'books', 'postgres')
64        cls.vl_authors = QgsVectorLayer(cls.dbconn + ' sslmode=disable key=\'pk\' table="qgis_test"."authors" sql=', 'authors', 'postgres')
65        cls.vl_editors = QgsVectorLayer(cls.dbconn + ' sslmode=disable key=\'fk_book,fk_author\' table="qgis_test"."editors" sql=', 'editors', 'postgres')
66        cls.vl_link_books_authors = QgsVectorLayer(cls.dbconn + ' sslmode=disable key=\'pk\' table="qgis_test"."books_authors" sql=', 'books_authors', 'postgres')
67
68        QgsProject.instance().addMapLayer(cls.vl_books)
69        QgsProject.instance().addMapLayer(cls.vl_authors)
70        QgsProject.instance().addMapLayer(cls.vl_editors)
71        QgsProject.instance().addMapLayer(cls.vl_link_books_authors)
72
73        cls.relMgr = QgsProject.instance().relationManager()
74
75        # Our mock QgsVectorLayerTools, that allow injecting data where user input is expected
76        cls.vltools = VlTools()
77        cls.layers = {cls.vl_authors, cls.vl_books, cls.vl_link_books_authors}
78
79        assert(cls.vl_authors.isValid())
80        assert(cls.vl_books.isValid())
81        assert(cls.vl_editors.isValid())
82        assert(cls.vl_link_books_authors.isValid())
83
84    @classmethod
85    def tearDownClass(cls):
86        QgsProject.instance().removeAllMapLayers()
87        cls.vl_books = None
88        cls.vl_authors = None
89        cls.vl_editors = None
90        cls.vl_link_books_authors = None
91        cls.layers = None
92        cls.mapCanvas = None
93        cls.vltools = None
94        cls.relMgr = None
95
96    def setUp(self):
97        self.rel_a = QgsRelation()
98        self.rel_a.setReferencingLayer(self.vl_link_books_authors.id())
99        self.rel_a.setReferencedLayer(self.vl_authors.id())
100        self.rel_a.addFieldPair('fk_author', 'pk')
101        self.rel_a.setId('rel_a')
102        assert(self.rel_a.isValid())
103        self.relMgr.addRelation(self.rel_a)
104
105        self.rel_b = QgsRelation()
106        self.rel_b.setReferencingLayer(self.vl_link_books_authors.id())
107        self.rel_b.setReferencedLayer(self.vl_books.id())
108        self.rel_b.addFieldPair('fk_book', 'pk')
109        self.rel_b.setId('rel_b')
110        assert(self.rel_b.isValid())
111        self.relMgr.addRelation(self.rel_b)
112
113        self.startTransaction()
114
115    def tearDown(self):
116        self.rollbackTransaction()
117        del self.transaction
118        self.relMgr.clear()
119
120    def startTransaction(self):
121        """
122        Start a new transaction and set all layers into transaction mode.
123
124        :return: None
125        """
126        self.transaction = QgsTransaction.create(self.layers)
127        self.transaction.begin()
128        for layer in self.layers:
129            layer.startEditing()
130
131    def rollbackTransaction(self):
132        """
133        Rollback all changes done in this transaction.
134        We always rollback and never commit to have the database in a pristine
135        state at the end of each test.
136
137        :return: None
138        """
139        for layer in self.layers:
140            layer.commitChanges()
141        self.transaction.rollback()
142
143    def test_delete_feature(self):
144        """
145        Check if a feature can be deleted properly
146        """
147        self.createWrapper(self.vl_authors, '"name"=\'Erich Gamma\'')
148
149        self.assertEqual(self.table_view.model().rowCount(), 1)
150
151        self.assertEqual(1, len([f for f in self.vl_books.getFeatures()]))
152
153        fid = next(self.vl_books.getFeatures(QgsFeatureRequest().setFilterExpression('"name"=\'Design Patterns. Elements of Reusable Object-Oriented Software\''))).id()
154
155        self.widget.featureSelectionManager().select([fid])
156
157        btn = self.widget.findChild(QToolButton, 'mDeleteFeatureButton')
158
159        def clickOk():
160            # Click the "Delete features" button on the confirmation message
161            # box
162            widget = self.widget.findChild(QMessageBox)
163            buttonBox = widget.findChild(QDialogButtonBox)
164            deleteButton = next((b for b in buttonBox.buttons() if buttonBox.buttonRole(b) == QDialogButtonBox.AcceptRole))
165            deleteButton.click()
166
167        QTimer.singleShot(1, clickOk)
168        btn.click()
169
170        # This is the important check that the feature is deleted
171        self.assertEqual(0, len([f for f in self.vl_books.getFeatures()]))
172
173        # This is actually more checking that the database on delete action is properly set on the relation
174        self.assertEqual(0, len([f for f in self.vl_link_books_authors.getFeatures()]))
175
176        self.assertEqual(self.table_view.model().rowCount(), 0)
177
178    def test_list(self):
179        """
180        Simple check if several related items are shown
181        """
182        wrapper = self.createWrapper(self.vl_books)  # NOQA
183
184        self.assertEqual(self.table_view.model().rowCount(), 4)
185
186    def test_add_feature(self):
187        """
188        Check if a new related feature is added
189        """
190        self.createWrapper(self.vl_authors, '"name"=\'Douglas Adams\'')
191
192        self.assertEqual(self.table_view.model().rowCount(), 0)
193
194        self.vltools.setValues([None, 'The Hitchhiker\'s Guide to the Galaxy', 'Sputnik Editions', 1961])
195        btn = self.widget.findChild(QToolButton, 'mAddFeatureButton')
196        btn.click()
197
198        # Book entry has been created
199        self.assertEqual(2, len([f for f in self.vl_books.getFeatures()]))
200
201        # Link entry has been created
202        self.assertEqual(5, len([f for f in self.vl_link_books_authors.getFeatures()]))
203
204        self.assertEqual(self.table_view.model().rowCount(), 1)
205
206    def test_link_feature(self):
207        """
208        Check if an existing feature can be linked
209        """
210        wrapper = self.createWrapper(self.vl_authors, '"name"=\'Douglas Adams\'')  # NOQA
211
212        f = QgsFeature(self.vl_books.fields())
213        f.setAttributes([self.vl_books.dataProvider().defaultValueClause(0), 'The Hitchhiker\'s Guide to the Galaxy', 'Sputnik Editions', 1961])
214        self.vl_books.addFeature(f)
215
216        btn = self.widget.findChild(QToolButton, 'mLinkFeatureButton')
217        btn.click()
218
219        dlg = self.widget.findChild(QDialog)
220        dlg.setSelectedFeatures([f.id()])
221        dlg.accept()
222
223        # magically the above code selects the feature here...
224        link_feature = next(self.vl_link_books_authors.getFeatures(QgsFeatureRequest().setFilterExpression('"fk_book"={}'.format(f[0]))))
225        self.assertIsNotNone(link_feature[0])
226
227        self.assertEqual(self.table_view.model().rowCount(), 1)
228
229    def test_unlink_feature(self):
230        """
231        Check if a linked feature can be unlinked
232        """
233        wrapper = self.createWrapper(self.vl_books)   # NOQA
234
235        # All authors are listed
236        self.assertEqual(self.table_view.model().rowCount(), 4)
237
238        it = self.vl_authors.getFeatures(
239            QgsFeatureRequest().setFilterExpression('"name" IN (\'Richard Helm\', \'Ralph Johnson\')'))
240
241        self.widget.featureSelectionManager().select([f.id() for f in it])
242
243        self.assertEqual(2, self.widget.featureSelectionManager().selectedFeatureCount())
244
245        btn = self.widget.findChild(QToolButton, 'mUnlinkFeatureButton')
246        btn.click()
247
248        # This is actually more checking that the database on delete action is properly set on the relation
249        self.assertEqual(2, len([f for f in self.vl_link_books_authors.getFeatures()]))
250
251        self.assertEqual(2, self.table_view.model().rowCount())
252
253    def test_discover_relations(self):
254        """
255        Test the automatic discovery of relations
256        """
257        relations = self.relMgr.discoverRelations([], [self.vl_authors, self.vl_books, self.vl_link_books_authors])
258        relations = {r.name(): r for r in relations}
259        self.assertEqual({'books_authors_fk_book_fkey', 'books_authors_fk_author_fkey'}, set(relations.keys()))
260
261        ba2b = relations['books_authors_fk_book_fkey']
262        self.assertTrue(ba2b.isValid())
263        self.assertEqual('books_authors', ba2b.referencingLayer().name())
264        self.assertEqual('books', ba2b.referencedLayer().name())
265        self.assertEqual([0], ba2b.referencingFields())
266        self.assertEqual([0], ba2b.referencedFields())
267
268        ba2a = relations['books_authors_fk_author_fkey']
269        self.assertTrue(ba2a.isValid())
270        self.assertEqual('books_authors', ba2a.referencingLayer().name())
271        self.assertEqual('authors', ba2a.referencedLayer().name())
272        self.assertEqual([1], ba2a.referencingFields())
273        self.assertEqual([0], ba2a.referencedFields())
274
275        self.assertEqual([], self.relMgr.discoverRelations([self.rel_a, self.rel_b], [self.vl_authors, self.vl_books, self.vl_link_books_authors]))
276        self.assertEqual(1, len(self.relMgr.discoverRelations([], [self.vl_authors, self.vl_link_books_authors])))
277
278        # composite keys relation
279        relations = self.relMgr.discoverRelations([], [self.vl_books, self.vl_editors])
280        self.assertEqual(len(relations), 1)
281        relation = relations[0]
282        self.assertEqual('books_fk_editor_fkey', relation.name())
283        self.assertTrue(relation.isValid())
284        self.assertEqual('books', relation.referencingLayer().name())
285        self.assertEqual('editors', relation.referencedLayer().name())
286        self.assertEqual([2, 3], relation.referencingFields())
287        self.assertEqual([0, 1], relation.referencedFields())
288
289    def test_selection(self):
290
291        fbook = QgsFeature(self.vl_books.fields())
292        fbook.setAttributes([self.vl_books.dataProvider().defaultValueClause(0), 'The Hitchhiker\'s Guide to the Galaxy', 'Sputnik Editions', 1961])
293        self.vl_books.addFeature(fbook)
294
295        flink = QgsFeature(self.vl_link_books_authors.fields())
296        flink.setAttributes([fbook.id(), 5])
297        self.vl_link_books_authors.addFeature(flink)
298
299        self.createWrapper(self.vl_authors, '"name"=\'Douglas Adams\'')
300
301        self.zoomToButton = self.widget.findChild(QToolButton, "mDeleteFeatureButton")
302        self.assertTrue(self.zoomToButton)
303        self.assertTrue(not self.zoomToButton.isEnabled())
304
305        selectionMgr = self.widget.featureSelectionManager()
306        self.assertTrue(selectionMgr)
307
308        self.vl_books.select(fbook.id())
309        self.assertEqual([fbook.id()], selectionMgr.selectedFeatureIds())
310        self.assertTrue(self.zoomToButton.isEnabled())
311
312        selectionMgr.deselect([fbook.id()])
313        self.assertEqual([], selectionMgr.selectedFeatureIds())
314        self.assertTrue(not self.zoomToButton.isEnabled())
315
316        self.vl_books.select([1, fbook.id()])
317        self.assertEqual([fbook.id()], selectionMgr.selectedFeatureIds())
318        self.assertTrue(self.zoomToButton.isEnabled())
319
320    def test_add_feature_geometry(self):
321        """
322        Test to add a feature with a geometry
323        """
324        vl_pipes = QgsVectorLayer(self.dbconn + ' sslmode=disable key=\'pk\' table="qgis_test"."pipes" (geom) sql=', 'pipes', 'postgres')
325        vl_leaks = QgsVectorLayer(self.dbconn + ' sslmode=disable key=\'pk\' table="qgis_test"."leaks" (geom) sql=', 'leaks', 'postgres')
326        vl_leaks.startEditing()
327
328        QgsProject.instance().addMapLayer(vl_pipes)
329        QgsProject.instance().addMapLayer(vl_leaks)
330
331        self.assertEqual(vl_pipes.featureCount(), 2)
332        self.assertEqual(vl_leaks.featureCount(), 3)
333
334        rel = QgsRelation()
335        rel.setReferencingLayer(vl_leaks.id())
336        rel.setReferencedLayer(vl_pipes.id())
337        rel.addFieldPair('pipe', 'id')
338        rel.setId('rel_pipe_leak')
339        self.assertTrue(rel.isValid())
340        self.relMgr.addRelation(rel)
341
342        # Mock vector layer tool to just set default value on created feature
343        class DummyVlTools(QgsVectorLayerTools):
344
345            def addFeature(self, layer, defaultValues, defaultGeometry):
346                f = QgsFeature(layer.fields())
347                for idx, value in defaultValues.items():
348                    f.setAttribute(idx, value)
349                f.setGeometry(defaultGeometry)
350                ok = layer.addFeature(f)
351
352                return ok, f
353
354        wrapper = QgsRelationWidgetWrapper(vl_leaks, rel)
355        context = QgsAttributeEditorContext()
356        vltool = DummyVlTools()
357        context.setVectorLayerTools(vltool)
358        context.setMapCanvas(self.mapCanvas)
359        cadDockWidget = QgsAdvancedDigitizingDockWidget(self.mapCanvas)
360        context.setCadDockWidget(cadDockWidget)
361        wrapper.setContext(context)
362        widget = wrapper.widget()
363        widget.show()
364        pipe = next(vl_pipes.getFeatures())
365        self.assertEqual(pipe.id(), 1)
366        wrapper.setFeature(pipe)
367        table_view = widget.findChild(QTableView)
368        self.assertEqual(table_view.model().rowCount(), 1)
369
370        btn = widget.findChild(QToolButton, 'mAddFeatureGeometryButton')
371        self.assertTrue(btn.isVisible())
372        self.assertTrue(btn.isEnabled())
373        btn.click()
374        self.assertTrue(self.mapCanvas.mapTool())
375        feature = QgsFeature(vl_leaks.fields())
376        feature.setGeometry(QgsGeometry.fromWkt('POINT(0 0.8)'))
377        self.mapCanvas.mapTool().digitizingCompleted.emit(feature)
378        self.assertEqual(table_view.model().rowCount(), 2)
379        self.assertEqual(vl_leaks.featureCount(), 4)
380        request = QgsFeatureRequest()
381        request.addOrderBy("id", False)
382
383        # get new created feature
384        feat = next(vl_leaks.getFeatures('"id" is NULL'))
385        self.assertTrue(feat.isValid())
386        self.assertTrue(feat.geometry().equals(QgsGeometry.fromWkt('POINT(0 0.8)')))
387
388        vl_leaks.rollBack()
389
390    def createWrapper(self, layer, filter=None):
391        """
392        Basic setup of a relation widget wrapper.
393        Will create a new wrapper and set its feature to the one and only book
394        in the table.
395        It will also assign some instance variables to help
396
397         * self.widget The created widget
398         * self.table_view The table view of the widget
399
400        :return: The created wrapper
401        """
402        if layer == self.vl_books:
403            relation = self.rel_b
404            nmrel = self.rel_a
405        else:
406            relation = self.rel_a
407            nmrel = self.rel_b
408
409        self.wrapper = QgsRelationWidgetWrapper(layer, relation)
410        self.wrapper.setConfig({'nm-rel': nmrel.id()})
411        context = QgsAttributeEditorContext()
412        context.setMapCanvas(self.mapCanvas)
413        context.setVectorLayerTools(self.vltools)
414        self.wrapper.setContext(context)
415
416        self.widget = self.wrapper.widget()
417        self.widget.show()
418
419        request = QgsFeatureRequest()
420        if filter:
421            request.setFilterExpression(filter)
422        book = next(layer.getFeatures(request))
423        self.wrapper.setFeature(book)
424
425        self.table_view = self.widget.findChild(QTableView)
426        return self.wrapper
427
428
429class VlTools(QgsVectorLayerTools):
430
431    """
432    Mock the QgsVectorLayerTools
433    Since we don't have a user on the test server to input this data for us, we can just use this.
434    """
435
436    def setValues(self, values):
437        """
438        Set the values for the next feature to insert
439        :param values: An array of values that shall be used for the next inserted record
440        :return: None
441        """
442        self.values = values
443
444    def addFeature(self, layer, defaultValues, defaultGeometry):
445        """
446        Overrides the addFeature method
447        :param layer: vector layer
448        :param defaultValues: some default values that may be provided by QGIS
449        :param defaultGeometry: a default geometry that may be provided by QGIS
450        :return: tuple(ok, f) where ok is if the layer added the feature and f is the added feature
451        """
452        values = list()
453        for i, v in enumerate(self.values):
454            if v:
455                values.append(v)
456            else:
457                values.append(layer.dataProvider().defaultValueClause(i))
458        f = QgsFeature(layer.fields())
459        f.setAttributes(self.values)
460        f.setGeometry(defaultGeometry)
461        ok = layer.addFeature(f)
462
463        return ok, f
464
465    def startEditing(self, layer):
466        pass
467
468    def stopEditing(self, layer, allowCancel):
469        pass
470
471    def saveEdits(self, layer):
472        pass
473
474
475if __name__ == '__main__':
476    unittest.main()
477