3 files changed, 28 insertions, 8 deletions
diff --git a/test/client.ts b/test/client.ts
new file mode 100644
index 0000000..182859b
--- /dev/null
+++ b/test/client.ts
@@ -0,0 +1,19 @@
+import { Schema } from 'mongoose';
+import { Client } from '../lib';
+import { Event as EventBase } from '../lib/event.model';
+import connection from './utils/dbConnection';
+
+interface Context {
+  message: string;
+}
+
+export type Event = EventBase<Context>;
+
+const contextSchema = new Schema({
+  message: String
+});
+
+const client = new Client<Context>(connection, contextSchema);
+
+export default client;
+
diff --git a/test/event.model.test.ts b/test/event.model.test.ts
index d7a4601..b06b268 100644
--- a/test/event.model.test.ts
+++ b/test/event.model.test.ts
@@ -1,10 +1,11 @@
 import { expect } from 'chai';
-import Model from './model';
-import connection from './utils/dbConnection';
+import client from './client';
+
+const Model = client.Event;
 
 describe('Event model', () => {
   after(async () => {
-    connection.dropCollection('customevents');
+    client.connection.dropCollection('events');
   });
 
   it('Should assign status and nextRunAt on creation', async () => {
diff --git a/test/scheduler.test.ts b/test/scheduler.test.ts
index ffcd623..1736ed4 100644
--- a/test/scheduler.test.ts
+++ b/test/scheduler.test.ts
@@ -1,12 +1,12 @@
 import { expect } from 'chai';
 import Scheduler from '../lib/scheduler';
-import Model, { CustomEvent } from './model';
-import connection from './utils/dbConnection';
+import client, { Event } from './client';
 
+const Model = client.Event;
 const sleep = async (time: number) => new Promise<void>(res => setTimeout(res, time));
 
 describe('Scheduler', async () => {
-  let event: CustomEvent;
+  let event: Event;
   let scheduler: Scheduler;
 
   beforeEach(() => {
@@ -28,14 +28,14 @@ describe('Scheduler', async () => {
   });
 
   after(async () => {
-    connection.dropCollection('customevents');
+    client.connection.dropCollection('events');
   });
 
 
   it('Should run event', async () => {
     // Wait until job is run
     await new Promise<void>(res => {
-      scheduler.registerHandler('test', async (event: CustomEvent) => {
+      scheduler.registerHandler('test', async (event: Event) => {
         await sleep(2000);
         await event.log(event.context.message)
         res();