11package org .springframework .tenancy .datasource ;
22
3- import java .sql .SQLException ;
4-
53import javax .persistence .EntityManager ;
64import javax .persistence .PersistenceContext ;
75
108import org .junit .Test ;
119import org .junit .runner .RunWith ;
1210import org .springframework .beans .factory .annotation .Autowired ;
13- import org .springframework .test .annotation .Rollback ;
11+ import org .springframework .test .annotation .DirtiesContext ;
1412import org .springframework .test .context .ContextConfiguration ;
1513import org .springframework .test .context .junit4 .SpringJUnit4ClassRunner ;
14+ import org .springframework .test .context .transaction .BeforeTransaction ;
15+ import org .springframework .transaction .PlatformTransactionManager ;
16+ import org .springframework .transaction .TransactionStatus ;
1617import org .springframework .transaction .annotation .Transactional ;
18+ import org .springframework .transaction .support .TransactionCallback ;
19+ import org .springframework .transaction .support .TransactionTemplate ;
1720
1821@ RunWith (SpringJUnit4ClassRunner .class )
19- @ ContextConfiguration ({ "/applicationContext-testDBSwitching.xml" })
20- @ Transactional ()
22+ @ ContextConfiguration
23+ @ Transactional
24+ @ DirtiesContext
2125public class DatabaseSwitchingDataSourceTest {
2226
2327 @ PersistenceContext
@@ -26,51 +30,48 @@ public class DatabaseSwitchingDataSourceTest {
2630 @ Autowired
2731 private BasicDataSource dataSource ;
2832
33+ @ Autowired
34+ private PlatformTransactionManager transactionManager ;
35+
2936 @ Autowired
3037 private ThreadlocalDatabaseSwitchingDataSource switchingDataSource ;
3138
39+ private static boolean initialized ;
40+
3241 private final String DB1 = "DB1" ;
3342 private final String DB2 = "DB2" ;
3443
3544 private final String CREATE_TABLE_STMT = "CREATE TABLE DOMAINOBJECT (ID INTEGER PRIMARY KEY, STRING CHAR(50)) GRANT ALL ON DOMAINOBJECT TO SA;" ;
3645
37- // I don't really want this to be a test, just a before that does not run under a transaction, however, I can't make
38- // that work (@BeforeTransaction). This means that we must run all tests in the class.
39- @ Rollback (false )
40- @ Test
41- public void setupDB1 () throws SQLException {
42- dataSource .getConnection ().createStatement ().execute ("create schema " + DB1 + " " + CREATE_TABLE_STMT );
46+ @ BeforeTransaction
47+ public void setUpDatabases () throws Exception {
4348
44- switchingDataSource .setDatabaseName (DB1 );
45- DomainObject do1 = getMockDomainObject ("db1" );
46- entityManager .persist (do1 );
47- DomainObject result = (DomainObject ) entityManager .createQuery ("select do from DomainObject do;" )
48- .getResultList ().get (0 );
49- Assert .assertNotNull (result );
50- Assert .assertEquals ("db1" , result .getString ());
51- entityManager .flush ();
49+ if (!initialized ) {
5250
53- }
51+ initialized = true ;
5452
55- // I don't really want this to be a test, just a before that does not run under a transaction, however, I can't make
56- // that work (@BeforeTransaction). This means that we must run all tests in the class.
57- @ Rollback (false )
58- @ Test
59- public void setupDB2 () throws SQLException {
60- dataSource .getConnection ().createStatement ().execute ("create schema " + DB2 + " " + CREATE_TABLE_STMT );
53+ dataSource .getConnection ().createStatement ().execute ("create schema " + DB1 + " " + CREATE_TABLE_STMT );
54+ dataSource .getConnection ().createStatement ().execute ("create schema " + DB2 + " " + CREATE_TABLE_STMT );
6155
62- switchingDataSource .setDatabaseName (DB2 );
63- entityManager .persist (getMockDomainObject ("db2" ));
64- DomainObject result = (DomainObject ) entityManager .createQuery ("select do from DomainObject do;" )
65- .getResultList ().get (0 );
66- Assert .assertNotNull (result );
67- Assert .assertEquals ("db2" , result .getString ());
68- entityManager .flush ();
56+ new TransactionTemplate (transactionManager ).execute (new TransactionCallback <Void >() {
57+ public Void doInTransaction (TransactionStatus status ) {
58+ setUpData (DB1 , "db1" );
59+ return null ;
60+ }
61+ });
62+ new TransactionTemplate (transactionManager ).execute (new TransactionCallback <Void >() {
63+ public Void doInTransaction (TransactionStatus status ) {
64+ setUpData (DB2 , "db2" );
65+ return null ;
66+ }
67+ });
68+
69+ }
6970
7071 }
7172
7273 @ Test
73- public void testGetFromD2 () {
74+ public void testGetFromDB2 () {
7475 switchingDataSource .setDatabaseName (DB2 );
7576 DomainObject result = (DomainObject ) entityManager .createQuery ("select do from DomainObject do;" )
7677 .getResultList ().get (0 );
@@ -89,6 +90,19 @@ public void testGetFromDB1() {
8990
9091 private static long nextId = 1 ;
9192
93+ private void setUpData (String databaseName , String objectName ) {
94+
95+ switchingDataSource .setDatabaseName (databaseName );
96+ DomainObject do1 = getMockDomainObject (objectName );
97+ entityManager .persist (do1 );
98+ DomainObject result = (DomainObject ) entityManager .createQuery ("select do from DomainObject do;" )
99+ .getResultList ().get (0 );
100+ Assert .assertNotNull (result );
101+ Assert .assertEquals (objectName , result .getString ());
102+ entityManager .flush ();
103+
104+ }
105+
92106 private DomainObject getMockDomainObject (String data ) {
93107 DomainObject result = new DomainObject ();
94108 result .setId (nextId ++);
0 commit comments