Error when loading in wagtailtrans own hooks

.coveragerc

@@ -4,4 +4,3 @@ omit =
     src/wagtailtrans/__init__.py
     src/wagtailtrans/admin.py
     src/wagtailtrans/config.py
-    src/wagtailtrans/wagtail_hooks.py

src/wagtailtrans/wagtail_hooks.py

@@ -41,13 +41,12 @@ def global_admin_js():
                 path=static('wagtailtrans/js/site_languages_editor.js'))
         )
 
-
 if not get_wagtailtrans_setting('SYNC_TREE'):
     """Only load hooks when WAGTAILTRANS_SYNC_TREE is disabled"""
 
     @hooks.register('register_page_listing_buttons')
     def page_translations_menu(page, page_perms, is_parent=False):
-        if not hasattr(page, 'language'):
+        if not issubclass(page.__class__, TranslatablePage):
             return
 
         if hasattr(page, 'canonical_page') and page.canonical_page:
@@ -116,7 +115,7 @@ def edit_in_language_button(page, page_perms, is_parent=False):
     clear interface to work in.
 
     """
-    if not hasattr(page, 'language'):
+    if not issubclass(page.__class__, TranslatablePage):
         return
 
     yield widgets.ButtonWithDropdownFromHook(

tests/_sandbox/pages/migrations/0003_article.py

@@ -0,0 +1,29 @@
+# Generated by Django 2.1.5 on 2019-02-12 08:23
+
+from django.db import migrations, models
+import django.db.models.deletion
+import wagtail.core.fields
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('wagtailcore', '0040_page_draft_title'),
+        ('pages', '0002_auto_20161214_0842'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Article',
+            fields=[
+                ('page_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, primary_key=True, serialize=False, to='wagtailcore.Page')),
+                ('date', models.DateField()),
+                ('intro_text', wagtail.core.fields.RichTextField()),
+                ('language', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, to='wagtailtrans.Language')),
+            ],
+            options={
+                'abstract': False,
+            },
+            bases=('wagtailcore.page',),
+        ),
+    ]

tests/_sandbox/pages/models.py

@@ -22,3 +22,16 @@ class HomePage(TranslatablePage):
     ]
 
     subpage_types = ['HomePage']
+
+
+class Article(Page):
+    date = models.DateField()
+    language = models.ForeignKey('wagtailtrans.Language', on_delete=models.SET_NULL, null=True)
+    intro_text = RichTextField()
+
+    content_panels = [
+        FieldPanel('title', 'full title'),
+        FieldPanel('date'),
+        FieldPanel('intro_text'),
+        FieldPanel('language'),
+    ]

tests/factories/pages.py

@@ -1,11 +1,14 @@
+import datetime
+
 import factory
+from factory.fuzzy import FuzzyDate
 from wagtail.core.models import Page
 from wagtail.images.tests.utils import (
     get_image_model, get_test_image_file)
 
 from wagtailtrans import models
 
-from tests._sandbox.pages.models import HomePage
+from tests._sandbox.pages.models import HomePage, Article
 from tests.factories import language
 
 
@@ -74,3 +77,35 @@ class WagtailPageFactory(factory.DjangoModelFactory):
 
     class Meta:
         model = Page
+
+    @classmethod
+    def _create(cls, *args, **kwargs):
+        model = args[0]
+        parent = kwargs.pop(
+            'parent',
+            Page.objects.get(path='0001')
+        )
+
+        try:
+            return model.objects.get(title=kwargs['title'])
+        except model.DoesNotExist:
+            kwargs['depth'] = parent.depth + 1
+
+            if kwargs['depth'] == 0:
+                return model.add_root(**kwargs)
+            else:
+                return parent.add_child(
+                    instance=model(
+                        **kwargs
+                    )
+                )
+
+
+class ArticleFactory(WagtailPageFactory):
+    title = 'Article 1'
+    intro_text = 'This is an article'
+    date = FuzzyDate(datetime.date(2019, 1, 1))
+    language = factory.SubFactory(language.LanguageFactory)
+
+    class Meta:
+        model = Article

tests/unit/test_wagtail_hooks.py

@@ -0,0 +1,50 @@
+import pytest
+
+from wagtail.admin import widgets
+from wagtail.core import hooks
+
+from tests.factories.pages import ArticleFactory
+from wagtailtrans.models import TranslatablePage
+from wagtailtrans.wagtail_hooks import edit_in_language_button
+
+
+@pytest.mark.django_db
+class TestWagtailHooks:
+    def setup(self):
+        self.article = ArticleFactory()
+        self.transpage = TranslatablePage()
+
+    @hooks.register('register_page_listing_buttons')
+    def page_translations_menu(self, page, page_perms, is_parent=False):
+        # To test this hook i need to add it here, if we don't do this
+        # we can't import the hook like the edit_in_language_button
+        if not issubclass(page.__class__, TranslatablePage):
+            return
+
+        if hasattr(page, 'canonical_page') and page.canonical_page:
+            return
+
+        yield widgets.ButtonWithDropdownFromHook(
+            'Translate into',
+            hook_name='wagtailtrans_dropdown_hook',
+            page=page,
+            page_perms=page_perms,
+            is_parent=is_parent,
+            priority=10
+        )
+
+    def test_edit_in_language_wagtail_hook_with_regular_page(self):
+        assert list(edit_in_language_button(self.article, page_perms=[])) == []
+
+    def test_edit_in_language_wagtail_hook_translateable_page(self):
+        result = list(edit_in_language_button(self.transpage, page_perms=[]))
+        assert len(result) == 1
+        assert result[0].label == 'Edit in'
+
+    def test_page_translations_menu_with_regular_page(self):
+        assert list(self.page_translations_menu(self.article, page_perms=[])) == []
+
+    def test_page_translations_menu_translateable_page(self):
+        result = list(self.page_translations_menu(self.transpage, page_perms=[]))
+        assert len(result) == 1
+        assert result[0].label == 'Translate into'