From 8a2dfb5b069106c9de2fa54af69caae80d4d1ebe Mon Sep 17 00:00:00 2001
From: Ivan Gonzalez <hi@dreamingecho.es>
Date: Tue, 1 Aug 2017 13:38:49 +0200
Subject: [PATCH] Add new count functions for multi tenancy in test support
 models

---
 test/support/simple_models.exs | 14 ++++++++++++++
 test/support/strict_models.exs | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/test/support/simple_models.exs b/test/support/simple_models.exs
index d1c5a400..c9846008 100644
--- a/test/support/simple_models.exs
+++ b/test/support/simple_models.exs
@@ -1,6 +1,8 @@
 defmodule SimpleCompany do
   use Ecto.Schema
 
+  alias PaperTrailTest.MultiTenantHelper, as: MultiTenant
+
   import Ecto.Changeset
   import Ecto.Query
 
@@ -31,11 +33,18 @@ defmodule SimpleCompany do
   def count do
     from(record in __MODULE__, select: count(record.id)) |> PaperTrail.RepoClient.repo.one
   end
+  def count(:multitenant) do
+    from(record in __MODULE__, select: count(record.id))
+    |> MultiTenant.add_prefix_to_query()
+    |> PaperTrail.RepoClient.repo.one
+  end
 end
 
 defmodule SimplePerson do
   use Ecto.Schema
 
+  alias PaperTrailTest.MultiTenantHelper, as: MultiTenant
+
   import Ecto.Changeset
   import Ecto.Query
 
@@ -62,4 +71,9 @@ defmodule SimplePerson do
   def count do
     from(record in __MODULE__, select: count(record.id)) |> PaperTrail.RepoClient.repo.one
   end
+  def count(:multitenant) do
+    from(record in __MODULE__, select: count(record.id))
+    |> MultiTenant.add_prefix_to_query()
+    |> PaperTrail.RepoClient.repo.one
+  end
 end
diff --git a/test/support/strict_models.exs b/test/support/strict_models.exs
index e7e24319..2217b994 100644
--- a/test/support/strict_models.exs
+++ b/test/support/strict_models.exs
@@ -1,6 +1,8 @@
 defmodule StrictCompany do
   use Ecto.Schema
 
+  alias PaperTrailTest.MultiTenantHelper, as: MultiTenant
+
   import Ecto.Changeset
   import Ecto.Query
 
@@ -34,11 +36,18 @@ defmodule StrictCompany do
   def count do
     from(record in __MODULE__, select: count(record.id)) |> PaperTrail.RepoClient.repo.one
   end
+  def count(:multitenant) do
+    from(record in __MODULE__, select: count(record.id))
+    |> MultiTenant.add_prefix_to_query()
+    |> PaperTrail.RepoClient.repo.one
+  end
 end
 
 defmodule StrictPerson do
   use Ecto.Schema
 
+  alias PaperTrailTest.MultiTenantHelper, as: MultiTenant
+
   import Ecto.Changeset
   import Ecto.Query
 
@@ -67,4 +76,9 @@ defmodule StrictPerson do
   def count do
     from(record in __MODULE__, select: count(record.id)) |> PaperTrail.RepoClient.repo.one
   end
+  def count(:multitenant) do
+    from(record in __MODULE__, select: count(record.id))
+    |> MultiTenant.add_prefix_to_query()
+    |> PaperTrail.RepoClient.repo.one
+  end
 end