diff --git a/src/extension.ts b/src/extension.ts
index 1c651a2..2f8d244 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -2,7 +2,7 @@
 // Import the module and reference it with the alias vscode in your code below
 import * as vscode from 'vscode';
 
-const CLASSNAME_REGEX = new RegExp(/className\s?=\s?['"](.*)['"]/);
+const CLASSNAME_REGEX = new RegExp(/className\s?=\s?['"]([^'"]*)['"]/);
 // This method is called when your extension is activated
 // Your extension is activated the very first time the command is executed
 export function activate(context: vscode.ExtensionContext) {
@@ -10,7 +10,7 @@ export function activate(context: vscode.ExtensionContext) {
 	// The command has been defined in the package.json file
 	// Now provide the implementation of the command with registerCommand
 	// The commandId parameter must match the command field in package.json
-	let disposable = vscode.commands.registerCommand('rmcss.modularize', () => {
+	let disposable = vscode.commands.registerCommand('rmcss.modularize', async () => {
 		// The code you place here will be executed every time your command is executed
 		// Display a message box to the user
 		const { activeTextEditor } = vscode.window;
@@ -20,7 +20,7 @@ export function activate(context: vscode.ExtensionContext) {
 		if (!activeTextEditor) { return; };
 		if(!selections?.length) { return; };
 		
-		activeTextEditor?.edit(editor => 
+		await activeTextEditor?.edit(editor => 
 			selections.forEach(selection => {
 				const finalSelection = document?.getWordRangeAtPosition(
 					selection.start,
@@ -30,7 +30,6 @@ export function activate(context: vscode.ExtensionContext) {
 				if (!finalSelection) { return; };
 
 				const text = activeTextEditor?.document.getText(finalSelection);
-
 				editor.replace(finalSelection, `className={${className}.${text.replace(CLASSNAME_REGEX, '$1')}}`);
 			})
 		);
diff --git a/src/test/suite/extension.test.ts b/src/test/suite/extension.test.ts
index a149b36..455db25 100644
--- a/src/test/suite/extension.test.ts
+++ b/src/test/suite/extension.test.ts
@@ -1,45 +1,114 @@
-import * as assert from 'assert';
-import * as vscode from 'vscode';
+import * as assert from "assert";
+import { afterEach } from "mocha";
+import * as vscode from "vscode";
 
-const SAMPLE = `
+suite("Extension Test Suite", () => {
+  vscode.window.showInformationMessage("Start all tests.");
+
+  afterEach(async () => {
+    await vscode.commands.executeCommand("workbench.action.closeActiveEditor");
+  });
+
+  test("Simple test", async () => {
+    const SAMPLE =
+`
 <div className="lala">LALA</div>
 <div className = 'lele'>LELE</div>
 <div className= "lili">LILI</div>
 <div something="whatever" className ="lulu">LULU</div>
 `;
 
-const EXPECTED = `
+    const EXPECTED =
+`
 <div className={styles.lala}>LALA</div>
 <div className={styles.lele}>LELE</div>
 <div className={styles.lili}>LILI</div>
 <div something="whatever" className={styles.lulu}>LULU</div>
 `;
+    // Create a new document with SAMPLE content
+    const document = await vscode.workspace.openTextDocument({
+      content: SAMPLE,
+    });
+
+    // Set document text editor as active
+    await vscode.window.showTextDocument(document);
+
+    // Put cursor on each className prop, value or =
+    if (!vscode.window.activeTextEditor) {
+      throw new Error();
+    }
+    vscode.window.activeTextEditor.selections = [
+      new vscode.Selection(1, 10, 1, 11),
+      new vscode.Selection(2, 19, 2, 19),
+      new vscode.Selection(3, 15, 3, 16),
+      new vscode.Selection(4, 30, 4, 30),
+    ];
+
+    // Run command
+    await vscode.commands.executeCommand("rmcss.modularize");
+
+    // Assert modular version is now the current text
+    assert.equal(document.getText(), EXPECTED);
+  });
+
+  test("Same line, single replace", async () => {
+    const SAMPLE =
+			`<div className="lala"> <div className = 'lele'>LELE</div> </div>`;
+
+    const EXPECTED =
+			`<div className={styles.lala}> <div className = 'lele'>LELE</div> </div>`;
+
+    // Create a new document with SAMPLE content
+    const document = await vscode.workspace.openTextDocument({
+      content: SAMPLE,
+    });
+
+    // Set document text editor as active
+    await vscode.window.showTextDocument(document);
+
+    // Put cursor on each className prop, value or =
+    if (!vscode.window.activeTextEditor) {
+      throw new Error();
+    }
+    vscode.window.activeTextEditor.selections = [
+      new vscode.Selection(0, 10, 0, 11),
+    ];
+
+    // Run command
+    await vscode.commands.executeCommand("rmcss.modularize");
+
+    // Assert modular version is now the current text
+    assert.equal(document.getText(), EXPECTED);
+  });
+
+  test("Same line, multiple replaces", async () => {
+    const SAMPLE =
+			`<div className="lala"> <div className = 'lele'>LELE</div> </div>`;
+
+    const EXPECTED =
+			`<div className={styles.lala}> <div className={styles.lele}>LELE</div> </div>`;
+
+    // Create a new document with SAMPLE content
+    const document = await vscode.workspace.openTextDocument({
+      content: SAMPLE,
+    });
+
+    // Set document text editor as active
+    await vscode.window.showTextDocument(document);
+
+    // Put cursor on each className prop, value or =
+    if (!vscode.window.activeTextEditor) {
+      throw new Error();
+    }
+    vscode.window.activeTextEditor.selections = [
+      new vscode.Selection(0, 10, 0, 11),
+      new vscode.Selection(0, 33, 0, 33),
+    ];
+
+    // Run command
+    await vscode.commands.executeCommand("rmcss.modularize");
 
-suite('Extension Test Suite', () => {
-	vscode.window.showInformationMessage('Start all tests.');
-
-	test('Simple test', async () => {
-		// Create a new document with SAMPLE content
-		const document = await vscode.workspace.openTextDocument({
-			content: SAMPLE
-		});
-
-		// Set document text editor as active
-		await vscode.window.showTextDocument(document);
-
-		// Put cursor on each className prop, value or =
-		if(!vscode.window.activeTextEditor) { return; };
-		vscode.window.activeTextEditor.selections = [
-			new vscode.Selection(1, 10, 1, 11),
-			new vscode.Selection(2, 19, 2, 19),
-			new vscode.Selection(3, 15, 3, 16),
-			new vscode.Selection(4, 30, 4, 30),
-		];
-
-		// Run command
-		await vscode.commands.executeCommand('rmcss.modularize');
-		
-		// Assert modular version is now the current text
-		assert.equal(document.getText(), EXPECTED);
-	});
+    // Assert modular version is now the current text
+    assert.equal(document.getText(), EXPECTED);
+  });
 });